<commit_before>package server

import (
	"errors"
	"fmt"
	"math"
	"math/big"
	"regexp"
	"strings"
	"sync"

	"github.com/golang/glog"

	ethcommon "github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/go-livepeer/core"
	"github.com/livepeer/go-livepeer/drivers"
	"github.com/livepeer/go-livepeer/monitor"
	"github.com/livepeer/go-livepeer/pm"

	"github.com/livepeer/lpms/stream"
)

type BroadcastSessionsManager struct {
	// Accessing or changing any of the below requires ownership of this mutex
	sessLock *sync.Mutex

	sessList []*BroadcastSession
	sessMap  map[string]*BroadcastSession
	numOrchs int // how many orchs to request at once

	refreshing bool // only allow one refresh in-flight
	finished   bool // set at stream end

	createSessions func() ([]*BroadcastSession, error)
}

func (bsm *BroadcastSessionsManager) selectSession() *BroadcastSession {
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()

	checkSessions := func(m *BroadcastSessionsManager) bool {
		numSess := len(m.sessList)
		if numSess < int(math.Ceil(float64(m.numOrchs)/2.0)) {
			go m.refreshSessions()
		}
		return numSess > 0
	}
	for checkSessions(bsm) {
		last := len(bsm.sessList) - 1
		sess, sessions := bsm.sessList[last], bsm.sessList[:last]
		bsm.sessList = sessions
		if _, ok := bsm.sessMap[sess.OrchestratorInfo.Transcoder]; ok {
			return sess
		}
		/*
		   Don't select sessions no longer in the map.

		   Retry if the first selected session has been removed from the map.
		   This may occur if the session is removed while still in the list.
		   To avoid a runtime search of the session list under lock, simply
		   fixup the session list at selection time by retrying the selection.
		*/
	}
	return nil
}

func (bsm *BroadcastSessionsManager) removeSession(session *BroadcastSession) {
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()
	delete(bsm.sessMap, session.OrchestratorInfo.Transcoder)
}

func (bsm *BroadcastSessionsManager) completeSession(sess *BroadcastSession) {
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()

	if _, ok := bsm.sessMap[sess.OrchestratorInfo.Transcoder]; ok {
		bsm.sessList = append(bsm.sessList, sess)
	}
}

func (bsm *BroadcastSessionsManager) refreshSessions() {

	glog.V(common.DEBUG).Info("Starting session refresh")
	defer glog.V(common.DEBUG).Info("Ending session refresh")
	bsm.sessLock.Lock()
	if bsm.finished || bsm.refreshing {
		bsm.sessLock.Unlock()
		return
	}
	bsm.refreshing = true
	bsm.sessLock.Unlock()

	newBroadcastSessions, err := bsm.createSessions()
	if err != nil {
		bsm.sessLock.Lock()
		bsm.refreshing = false
		bsm.sessLock.Unlock()
		return
	}

	// if newBroadcastSessions is empty, exit without refreshing list
	if len(newBroadcastSessions) <= 0 {
		bsm.sessLock.Lock()
		bsm.refreshing = false
		bsm.sessLock.Unlock()
		return
	}

	uniqueSessions := make([]*BroadcastSession, 0, len(newBroadcastSessions))
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()

	bsm.refreshing = false
	if bsm.finished {
		return
	}

	for _, sess := range newBroadcastSessions {
		if _, ok := bsm.sessMap[sess.OrchestratorInfo.Transcoder]; ok {
			continue
		}
		uniqueSessions = append(uniqueSessions, sess)
		bsm.sessMap[sess.OrchestratorInfo.Transcoder] = sess
	}
	bsm.sessList = append(uniqueSessions, bsm.sessList...)
}

func (bsm *BroadcastSessionsManager) cleanup() {
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()
	bsm.finished = true
	bsm.sessList = nil
	bsm.sessMap = make(map[string]*BroadcastSession) // prevent segfaults
}

func NewSessionManager(node *core.LivepeerNode, pl core.PlaylistManager) *BroadcastSessionsManager {
	var poolSize float64
	if node.OrchestratorPool != nil {
		poolSize = float64(node.OrchestratorPool.Size())
	}
	maxInflight := HTTPTimeout.Seconds() / SegLen.Seconds()
	numOrchs := int(math.Min(poolSize, maxInflight*2))
	bsm := &BroadcastSessionsManager{
		sessMap:        make(map[string]*BroadcastSession),
		createSessions: func() ([]*BroadcastSession, error) { return selectOrchestrator(node, pl, numOrchs) },
		sessLock:       &sync.Mutex{},
		numOrchs:       numOrchs,
	}
	bsm.refreshSessions()
	return bsm
}

func selectOrchestrator(n *core.LivepeerNode, cpl core.PlaylistManager, count int) ([]*BroadcastSession, error) {
	if n.OrchestratorPool == nil {
		glog.Info("No orchestrators specified; not transcoding")
		return nil, ErrDiscovery
	}

	rpcBcast := core.NewBroadcaster(n)

	tinfos, err := n.OrchestratorPool.GetOrchestrators(count)
	if len(tinfos) <= 0 {
		glog.Info("No orchestrators found; not transcoding. Error: ", err)
		return nil, ErrNoOrchs
	}
	if err != nil {
		return nil, err
	}

	var sessions []*BroadcastSession

	for _, tinfo := range tinfos {
		var sessionID string

		if n.Sender != nil {
			protoParams := tinfo.TicketParams
			params := pm.TicketParams{
				Recipient:         ethcommon.BytesToAddress(protoParams.Recipient),
				FaceValue:         new(big.Int).SetBytes(protoParams.FaceValue),
				WinProb:           new(big.Int).SetBytes(protoParams.WinProb),
				RecipientRandHash: ethcommon.BytesToHash(protoParams.RecipientRandHash),
				Seed:              new(big.Int).SetBytes(protoParams.Seed),
			}

			sessionID = n.Sender.StartSession(params)
		}

		var orchOS drivers.OSSession
		if len(tinfo.Storage) > 0 {
			orchOS = drivers.NewSession(tinfo.Storage[0])
		}

		bcastOS := cpl.GetOSSession()
		if bcastOS.IsExternal() {
			// Give each O its own OS session to prevent front running uploads
			pfx := fmt.Sprintf("%v/%v", cpl.ManifestID(), core.RandomManifestID())
			bcastOS = drivers.NodeStorage.NewSession(pfx)
		}

		session := &BroadcastSession{
			Broadcaster:      rpcBcast,
			ManifestID:       cpl.ManifestID(),
			Profiles:         BroadcastJobVideoProfiles,
			OrchestratorInfo: tinfo,
			OrchestratorOS:   orchOS,
			BroadcasterOS:    bcastOS,
			Sender:           n.Sender,
			PMSessionID:      sessionID,
		}

		sessions = append(sessions, session)
	}
	return sessions, nil
}

func processSegment(cxn *rtmpConnection, seg *stream.HLSSegment) {

	nonce := cxn.nonce
	rtmpStrm := cxn.stream
	cpl := cxn.pl
	mid := cxn.mid
	vProfile := cxn.profile

	if monitor.Enabled {
		monitor.LogSegmentEmerged(nonce, seg.SeqNo, len(BroadcastJobVideoProfiles))
	}

	seg.Name = "" // hijack seg.Name to convey the uploaded URI
	name := fmt.Sprintf("%s/%d.ts", vProfile.Name, seg.SeqNo)
	uri, err := cpl.GetOSSession().SaveData(name, seg.Data)
	if err != nil {
		glog.Errorf("Error saving segment %d: %v", seg.SeqNo, err)
		if monitor.Enabled {
			monitor.LogSegmentUploadFailed(nonce, seg.SeqNo, monitor.SegmentUploadErrorUnknown, err.Error())
		}
		return
	}
	if cpl.GetOSSession().IsExternal() {
		seg.Name = uri // hijack seg.Name to convey the uploaded URI
	}
	err = cpl.InsertHLSSegment(vProfile, seg.SeqNo, uri, seg.Duration)
	if monitor.Enabled {
		monitor.LogSourceSegmentAppeared(nonce, seg.SeqNo, string(mid), vProfile.Name)
		glog.V(6).Infof("Appeared segment %d", seg.SeqNo)
	}
	if err != nil {
		glog.Errorf("Error inserting segment %d: %v", seg.SeqNo, err)
		if monitor.Enabled {
			monitor.LogSegmentUploadFailed(nonce, seg.SeqNo, monitor.SegmentUploadErrorUnknown, err.Error())
		}
	}

	sess := cxn.sessManager.selectSession()
	// Return early under a few circumstances:
	// View-only (non-transcoded) streams or mid-failover
	if sess == nil {
		if monitor.Enabled {
			monitor.LogSegmentTranscodeFailed(monitor.SegmentTranscodeErrorNoOrchestrators, nonce, seg.SeqNo, errors.New("No Orchestrators Error"))
		}
		return
	}

	// Process the rest of the segment asynchronously - transcode
	go func() {
		// storage the orchestrator prefers
		if ios := sess.OrchestratorOS; ios != nil {
			// XXX handle case when orch expects direct upload
			uri, err := ios.SaveData(name, seg.Data)
			if err != nil {
				glog.Error("Error saving segment to OS ", err)
				if monitor.Enabled {
					monitor.LogSegmentUploadFailed(nonce, seg.SeqNo, monitor.SegmentUploadErrorOS, err.Error())
				}
				cxn.sessManager.completeSession(sess)
				return
			}
			seg.Name = uri // hijack seg.Name to convey the uploaded URI
		}

		// send segment to the orchestrator
		glog.V(common.DEBUG).Infof("Submitting segment %d", seg.SeqNo)

		res, err := SubmitSegment(sess, seg, nonce)
		if err != nil {
			if shouldStopStream(err) {
				glog.Warningf("Stopping current stream due to: %v", err)
				rtmpStrm.Close()
				return
			}
			if shouldStopSession(err) {
				cxn.sessManager.removeSession(sess)
			}
			return
		}

		cxn.sessManager.completeSession(sess)

		// download transcoded segments from the transcoder
		gotErr := false // only send one error msg per segment list
		var errCode monitor.SegmentTranscodeError
		errFunc := func(subType monitor.SegmentTranscodeError, url string, err error) {
			glog.Errorf("%v error with segment %v: %v (URL: %v)", subType, seg.SeqNo, err, url)
			if monitor.Enabled && !gotErr {
				monitor.LogSegmentTranscodeFailed(subType, nonce, seg.SeqNo, err)
				gotErr = true
				errCode = subType
			}
		}
		if res == nil {
			return
		}

		segHashes := make([][]byte, len(res.Segments))
		n := len(res.Segments)
		segHashLock := &sync.Mutex{}
		cond := sync.NewCond(segHashLock)

		dlFunc := func(url string, i int) {
			defer func() {
				cond.L.Lock()
				n--
				if n == 0 {
					cond.Signal()
				}
				cond.L.Unlock()
			}()

			if bos := sess.BroadcasterOS; bos != nil && !drivers.IsOwnExternal(url) {
				data, err := drivers.GetSegmentData(url)
				if err != nil {
					errFunc(monitor.SegmentTranscodeErrorDownload, url, err)
					return
				}
				name := fmt.Sprintf("%s/%d.ts", sess.Profiles[i].Name, seg.SeqNo)
				newUrl, err := bos.SaveData(name, data)
				if err != nil {
					switch err.Error() {
					case "Session ended":
						errFunc(monitor.SegmentTranscodeErrorSessionEnded, url, err)
					default:
						errFunc(monitor.SegmentTranscodeErrorSaveData, url, err)
					}
					return
				}
				url = newUrl

				hash := crypto.Keccak256(data)
				segHashLock.Lock()
				segHashes[i] = hash
				segHashLock.Unlock()
			}

			if monitor.Enabled {
				monitor.LogTranscodedSegmentAppeared(nonce, seg.SeqNo, sess.Profiles[i].Name)
			}
			err = cpl.InsertHLSSegment(&sess.Profiles[i], seg.SeqNo, url, seg.Duration)
			if err != nil {
				errFunc(monitor.SegmentTranscodeErrorPlaylist, url, err)
				return
			}
		}

		for i, v := range res.Segments {
			go dlFunc(v.Url, i)
		}

		cond.L.Lock()
		for n != 0 {
			cond.Wait()
		}
		cond.L.Unlock()
		if monitor.Enabled {
			monitor.SegmentFullyTranscoded(nonce, seg.SeqNo, common.ProfilesNames(sess.Profiles), errCode)
		}

		ticketParams := sess.OrchestratorInfo.GetTicketParams()
		if ticketParams != nil && // may be nil in offchain mode
			!pm.VerifySig(ethcommon.BytesToAddress(ticketParams.Recipient), crypto.Keccak256(segHashes...), res.Sig) {
			glog.Error("Sig check failed for segment ", seg.SeqNo)
			return
		}

		glog.V(common.DEBUG).Info("Successfully validated segment ", seg.SeqNo)
	}()
}

var sessionErrStrings = []string{"dial tcp", "unexpected EOF", core.ErrOrchBusy.Error(), core.ErrOrchCap.Error()}

func generateSessionErrors() *regexp.Regexp {
	// Given a list [err1, err2, err3] generates a regexp `(err1)|(err2)|(err3)`
	groups := []string{}
	for _, v := range sessionErrStrings {
		groups = append(groups, fmt.Sprintf("(%v)", v))
	}
	return regexp.MustCompile(strings.Join(groups, "|"))
}

var sessionErrRegex = generateSessionErrors()

func shouldStopSession(err error) bool {
	return sessionErrRegex.MatchString(err.Error())
}
<commit_msg>broadcast: Add rudimentary retries.<commit_after>package server

import (
	"errors"
	"fmt"
	"math"
	"math/big"
	"regexp"
	"strings"
	"sync"

	"github.com/golang/glog"

	ethcommon "github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/go-livepeer/core"
	"github.com/livepeer/go-livepeer/drivers"
	"github.com/livepeer/go-livepeer/monitor"
	"github.com/livepeer/go-livepeer/pm"

	"github.com/livepeer/lpms/stream"
)

type BroadcastSessionsManager struct {
	// Accessing or changing any of the below requires ownership of this mutex
	sessLock *sync.Mutex

	sessList []*BroadcastSession
	sessMap  map[string]*BroadcastSession
	numOrchs int // how many orchs to request at once

	refreshing bool // only allow one refresh in-flight
	finished   bool // set at stream end

	createSessions func() ([]*BroadcastSession, error)
}

func (bsm *BroadcastSessionsManager) selectSession() *BroadcastSession {
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()

	checkSessions := func(m *BroadcastSessionsManager) bool {
		numSess := len(m.sessList)
		if numSess < int(math.Ceil(float64(m.numOrchs)/2.0)) {
			go m.refreshSessions()
		}
		return numSess > 0
	}
	for checkSessions(bsm) {
		last := len(bsm.sessList) - 1
		sess, sessions := bsm.sessList[last], bsm.sessList[:last]
		bsm.sessList = sessions
		if _, ok := bsm.sessMap[sess.OrchestratorInfo.Transcoder]; ok {
			return sess
		}
		/*
		   Don't select sessions no longer in the map.

		   Retry if the first selected session has been removed from the map.
		   This may occur if the session is removed while still in the list.
		   To avoid a runtime search of the session list under lock, simply
		   fixup the session list at selection time by retrying the selection.
		*/
	}
	return nil
}

func (bsm *BroadcastSessionsManager) removeSession(session *BroadcastSession) {
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()
	delete(bsm.sessMap, session.OrchestratorInfo.Transcoder)
}

func (bsm *BroadcastSessionsManager) completeSession(sess *BroadcastSession) {
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()

	if _, ok := bsm.sessMap[sess.OrchestratorInfo.Transcoder]; ok {
		bsm.sessList = append(bsm.sessList, sess)
	}
}

func (bsm *BroadcastSessionsManager) refreshSessions() {

	glog.V(common.DEBUG).Info("Starting session refresh")
	defer glog.V(common.DEBUG).Info("Ending session refresh")
	bsm.sessLock.Lock()
	if bsm.finished || bsm.refreshing {
		bsm.sessLock.Unlock()
		return
	}
	bsm.refreshing = true
	bsm.sessLock.Unlock()

	newBroadcastSessions, err := bsm.createSessions()
	if err != nil {
		bsm.sessLock.Lock()
		bsm.refreshing = false
		bsm.sessLock.Unlock()
		return
	}

	// if newBroadcastSessions is empty, exit without refreshing list
	if len(newBroadcastSessions) <= 0 {
		bsm.sessLock.Lock()
		bsm.refreshing = false
		bsm.sessLock.Unlock()
		return
	}

	uniqueSessions := make([]*BroadcastSession, 0, len(newBroadcastSessions))
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()

	bsm.refreshing = false
	if bsm.finished {
		return
	}

	for _, sess := range newBroadcastSessions {
		if _, ok := bsm.sessMap[sess.OrchestratorInfo.Transcoder]; ok {
			continue
		}
		uniqueSessions = append(uniqueSessions, sess)
		bsm.sessMap[sess.OrchestratorInfo.Transcoder] = sess
	}
	bsm.sessList = append(uniqueSessions, bsm.sessList...)
}

func (bsm *BroadcastSessionsManager) cleanup() {
	bsm.sessLock.Lock()
	defer bsm.sessLock.Unlock()
	bsm.finished = true
	bsm.sessList = nil
	bsm.sessMap = make(map[string]*BroadcastSession) // prevent segfaults
}

func NewSessionManager(node *core.LivepeerNode, pl core.PlaylistManager) *BroadcastSessionsManager {
	var poolSize float64
	if node.OrchestratorPool != nil {
		poolSize = float64(node.OrchestratorPool.Size())
	}
	maxInflight := HTTPTimeout.Seconds() / SegLen.Seconds()
	numOrchs := int(math.Min(poolSize, maxInflight*2))
	bsm := &BroadcastSessionsManager{
		sessMap:        make(map[string]*BroadcastSession),
		createSessions: func() ([]*BroadcastSession, error) { return selectOrchestrator(node, pl, numOrchs) },
		sessLock:       &sync.Mutex{},
		numOrchs:       numOrchs,
	}
	bsm.refreshSessions()
	return bsm
}

func selectOrchestrator(n *core.LivepeerNode, cpl core.PlaylistManager, count int) ([]*BroadcastSession, error) {
	if n.OrchestratorPool == nil {
		glog.Info("No orchestrators specified; not transcoding")
		return nil, ErrDiscovery
	}

	rpcBcast := core.NewBroadcaster(n)

	tinfos, err := n.OrchestratorPool.GetOrchestrators(count)
	if len(tinfos) <= 0 {
		glog.Info("No orchestrators found; not transcoding. Error: ", err)
		return nil, ErrNoOrchs
	}
	if err != nil {
		return nil, err
	}

	var sessions []*BroadcastSession

	for _, tinfo := range tinfos {
		var sessionID string

		if n.Sender != nil {
			protoParams := tinfo.TicketParams
			params := pm.TicketParams{
				Recipient:         ethcommon.BytesToAddress(protoParams.Recipient),
				FaceValue:         new(big.Int).SetBytes(protoParams.FaceValue),
				WinProb:           new(big.Int).SetBytes(protoParams.WinProb),
				RecipientRandHash: ethcommon.BytesToHash(protoParams.RecipientRandHash),
				Seed:              new(big.Int).SetBytes(protoParams.Seed),
			}

			sessionID = n.Sender.StartSession(params)
		}

		var orchOS drivers.OSSession
		if len(tinfo.Storage) > 0 {
			orchOS = drivers.NewSession(tinfo.Storage[0])
		}

		bcastOS := cpl.GetOSSession()
		if bcastOS.IsExternal() {
			// Give each O its own OS session to prevent front running uploads
			pfx := fmt.Sprintf("%v/%v", cpl.ManifestID(), core.RandomManifestID())
			bcastOS = drivers.NodeStorage.NewSession(pfx)
		}

		session := &BroadcastSession{
			Broadcaster:      rpcBcast,
			ManifestID:       cpl.ManifestID(),
			Profiles:         BroadcastJobVideoProfiles,
			OrchestratorInfo: tinfo,
			OrchestratorOS:   orchOS,
			BroadcasterOS:    bcastOS,
			Sender:           n.Sender,
			PMSessionID:      sessionID,
		}

		sessions = append(sessions, session)
	}
	return sessions, nil
}

func processSegment(cxn *rtmpConnection, seg *stream.HLSSegment) {

	nonce := cxn.nonce
	cpl := cxn.pl
	mid := cxn.mid
	vProfile := cxn.profile

	glog.V(common.DEBUG).Infof("Processing segment %d", seg.SeqNo)
	if monitor.Enabled {
		monitor.LogSegmentEmerged(nonce, seg.SeqNo, len(BroadcastJobVideoProfiles))
	}

	seg.Name = "" // hijack seg.Name to convey the uploaded URI
	name := fmt.Sprintf("%s/%d.ts", vProfile.Name, seg.SeqNo)
	uri, err := cpl.GetOSSession().SaveData(name, seg.Data)
	if err != nil {
		glog.Errorf("Error saving segment %d: %v", seg.SeqNo, err)
		if monitor.Enabled {
			monitor.LogSegmentUploadFailed(nonce, seg.SeqNo, monitor.SegmentUploadErrorUnknown, err.Error())
		}
		return
	}
	if cpl.GetOSSession().IsExternal() {
		seg.Name = uri // hijack seg.Name to convey the uploaded URI
	}
	err = cpl.InsertHLSSegment(vProfile, seg.SeqNo, uri, seg.Duration)
	if monitor.Enabled {
		monitor.LogSourceSegmentAppeared(nonce, seg.SeqNo, string(mid), vProfile.Name)
	}
	if err != nil {
		glog.Errorf("Error inserting segment %d: %v", seg.SeqNo, err)
		if monitor.Enabled {
			monitor.LogSegmentUploadFailed(nonce, seg.SeqNo, monitor.SegmentUploadErrorUnknown, err.Error())
		}
	}

	// Process the rest of the segment asynchronously - transcode
	go func() {
		for true {
			// if fails, retry; rudimentary
			if err := transcodeSegment(cxn, seg, name); err == nil {
				return
			}
		}
	}()
}

func transcodeSegment(cxn *rtmpConnection, seg *stream.HLSSegment, name string) error {

	nonce := cxn.nonce
	rtmpStrm := cxn.stream
	cpl := cxn.pl
	sess := cxn.sessManager.selectSession()
	// Return early under a few circumstances:
	// View-only (non-transcoded) streams or no sessions available
	if sess == nil {
		if monitor.Enabled {
			monitor.LogSegmentTranscodeFailed(monitor.SegmentTranscodeErrorNoOrchestrators, nonce, seg.SeqNo, errors.New("No Orchestrators Error"))
		}
		glog.Info("No sessions available for segment ", seg.SeqNo)
		return nil
	}
	{

		// storage the orchestrator prefers
		if ios := sess.OrchestratorOS; ios != nil {
			// XXX handle case when orch expects direct upload
			uri, err := ios.SaveData(name, seg.Data)
			if err != nil {
				glog.Error("Error saving segment to OS ", err)
				if monitor.Enabled {
					monitor.LogSegmentUploadFailed(nonce, seg.SeqNo, monitor.SegmentUploadErrorOS, err.Error())
				}
				cxn.sessManager.removeSession(sess)
				return err
			}
			seg.Name = uri // hijack seg.Name to convey the uploaded URI
		}

		// send segment to the orchestrator
		glog.V(common.DEBUG).Infof("Submitting segment %d", seg.SeqNo)

		res, err := SubmitSegment(sess, seg, nonce)
		if err != nil || res == nil {
			cxn.sessManager.removeSession(sess)
			if res == nil && err == nil {
				return errors.New("Empty response")
			}
			if shouldStopStream(err) {
				glog.Warningf("Stopping current stream due to: %v", err)
				rtmpStrm.Close()
				return err
			}
			if shouldStopSession(err) {
			}
			return err
		}

		cxn.sessManager.completeSession(sess)

		// download transcoded segments from the transcoder
		gotErr := false // only send one error msg per segment list
		var errCode monitor.SegmentTranscodeError
		errFunc := func(subType monitor.SegmentTranscodeError, url string, err error) {
			glog.Errorf("%v error with segment %v: %v (URL: %v)", subType, seg.SeqNo, err, url)
			if monitor.Enabled && !gotErr {
				monitor.LogSegmentTranscodeFailed(subType, nonce, seg.SeqNo, err)
				gotErr = true
				errCode = subType
			}
		}

		segHashes := make([][]byte, len(res.Segments))
		n := len(res.Segments)
		segHashLock := &sync.Mutex{}
		cond := sync.NewCond(segHashLock)

		dlFunc := func(url string, i int) {
			defer func() {
				cond.L.Lock()
				n--
				if n == 0 {
					cond.Signal()
				}
				cond.L.Unlock()
			}()

			if bos := sess.BroadcasterOS; bos != nil && !drivers.IsOwnExternal(url) {
				data, err := drivers.GetSegmentData(url)
				if err != nil {
					errFunc(monitor.SegmentTranscodeErrorDownload, url, err)
					return
				}
				name := fmt.Sprintf("%s/%d.ts", sess.Profiles[i].Name, seg.SeqNo)
				newUrl, err := bos.SaveData(name, data)
				if err != nil {
					switch err.Error() {
					case "Session ended":
						errFunc(monitor.SegmentTranscodeErrorSessionEnded, url, err)
					default:
						errFunc(monitor.SegmentTranscodeErrorSaveData, url, err)
					}
					return
				}
				url = newUrl

				hash := crypto.Keccak256(data)
				segHashLock.Lock()
				segHashes[i] = hash
				segHashLock.Unlock()
			}

			if monitor.Enabled {
				monitor.LogTranscodedSegmentAppeared(nonce, seg.SeqNo, sess.Profiles[i].Name)
			}
			err = cpl.InsertHLSSegment(&sess.Profiles[i], seg.SeqNo, url, seg.Duration)
			if err != nil {
				errFunc(monitor.SegmentTranscodeErrorPlaylist, url, err)
				return
			}
		}

		for i, v := range res.Segments {
			go dlFunc(v.Url, i)
		}

		cond.L.Lock()
		for n != 0 {
			cond.Wait()
		}
		cond.L.Unlock()
		if monitor.Enabled {
			monitor.SegmentFullyTranscoded(nonce, seg.SeqNo, common.ProfilesNames(sess.Profiles), errCode)
		}

		ticketParams := sess.OrchestratorInfo.GetTicketParams()
		if ticketParams != nil && // may be nil in offchain mode
			!pm.VerifySig(ethcommon.BytesToAddress(ticketParams.Recipient), crypto.Keccak256(segHashes...), res.Sig) {
			glog.Error("Sig check failed for segment ", seg.SeqNo)
			return errors.New("PM Check Failed")
		}

		glog.V(common.DEBUG).Info("Successfully validated segment ", seg.SeqNo)
		return nil
	}
}

var sessionErrStrings = []string{"dial tcp", "unexpected EOF", core.ErrOrchBusy.Error(), core.ErrOrchCap.Error()}

func generateSessionErrors() *regexp.Regexp {
	// Given a list [err1, err2, err3] generates a regexp `(err1)|(err2)|(err3)`
	groups := []string{}
	for _, v := range sessionErrStrings {
		groups = append(groups, fmt.Sprintf("(%v)", v))
	}
	return regexp.MustCompile(strings.Join(groups, "|"))
}

var sessionErrRegex = generateSessionErrors()

func shouldStopSession(err error) bool {
	return sessionErrRegex.MatchString(err.Error())
}
<|endoftext|>


<commit_before>package ops

import (
	"errors"
	"github.com/phzfi/RIC/server/logging"
)

var (
	ErrRootNotFound     = errors.New("Root not found")
	ErrRootAlreadyAdded = errors.New("Root is already served")
)

type roots []string

func (roots *roots) Add(n string) error {
	logging.Debug("Adding root: " + n)
	for _, path := range *roots {
		if path == n {
			return ErrRootAlreadyAdded
		}
	}

	*roots = append(*roots, n)
	return nil
}

func (roots *roots) Remove(r string) error {
	for i, path := range *roots {
		if path == r {
			*roots = append((*roots)[:i], (*roots)[i+1:]...)
			return nil
		}
	}

	return ErrRootNotFound
}
<commit_msg>make checking whether a root exists a function<commit_after>package ops

import (
	"errors"
	"github.com/phzfi/RIC/server/logging"
)

var (
	ErrRootNotFound     = errors.New("Root not found")
	ErrRootAlreadyAdded = errors.New("Root is already served")
)

type roots []string

func (roots *roots) Add(r string) error {
	logging.Debug("Adding root: " + r)
	if roots.HasRoot(r) {
		return ErrRootAlreadyAdded
	}

	*roots = append(*roots, r)
	return nil
}

func (roots *roots) HasRoot(r string) bool {
	for _, path := range *roots {
		if path == r {
			return true
		}
	}
	return false
}

func (roots *roots) Remove(r string) error {
	for i, path := range *roots {
		if path == r {
			*roots = append((*roots)[:i], (*roots)[i+1:]...)
			return nil
		}
	}

	return ErrRootNotFound
}
<|endoftext|>


<commit_before>package web

// import (
// 	"bytes"
// 	"encoding/json"
// 	"fmt"
// 	"io"
// 	"io/ioutil"
// 	"net/http/httptest"
// 	"os"
// 	"path"
// 	"reflect"
// 	"testing"

// 	"github.com/gin-gonic/gin"

// 	"github.com/bobinette/papernet"
// 	"github.com/bobinette/papernet/auth"
// 	"github.com/bobinette/papernet/bleve"
// 	"github.com/bobinette/papernet/bolt"
// )

// type handlers struct {
// 	PaperHandler *PaperHandler
// 	UserHandler  *UserHandler
// }

// func createRouter(t *testing.T) (*gin.Engine, handlers, func()) {
// 	dir, err := ioutil.TempDir("", "")
// 	if err != nil {
// 		t.Fatal("could not create tmp file:", err)
// 	}

// 	index := &bleve.PaperIndex{}
// 	err = index.Open(path.Join(dir, "test", "index"))
// 	if err != nil {
// 		t.Fatal("error creating index", err)
// 	}

// 	driver := &bolt.Driver{}
// 	err = driver.Open(path.Join(dir, "test", "db"))
// 	if err != nil {
// 		t.Fatal("error creating index", err)
// 	}

// 	encoder := auth.Encoder{Key: "test"}

// 	paperStore := &bolt.PaperStore{Driver: driver}
// 	userRepo := &bolt.UserRepository{Driver: driver}
// 	tagIndex := &bolt.TagIndex{Driver: driver}

// 	paperHandler := &PaperHandler{
// 		Store:          paperStore,
// 		Searcher:       index,
// 		TagIndex:       tagIndex,
// 		UserRepository: userRepo,
// 		Authenticator: Authenticator{
// 			UserRepository: userRepo,
// 			Encoder:        encoder,
// 		},
// 	}

// 	userHandler := &UserHandler{
// 		Authenticator: Authenticator{
// 			UserRepository: userRepo,
// 			Encoder:        encoder,
// 		},
// 		Repository: userRepo,
// 	}

// 	gin.SetMode(gin.ReleaseMode) // avoid unnecessary log
// 	router := gin.New()
// 	paperHandler.RegisterRoutes(router)
// 	userHandler.RegisterRoutes(router)

// 	return router, handlers{
// 			PaperHandler: paperHandler,
// 			UserHandler:  userHandler,
// 		}, func() {
// 			if err := index.Close(); err != nil {
// 				t.Log(err)
// 			}
// 			if err := os.RemoveAll(dir); err != nil {
// 				t.Log(err)
// 			}
// 		}
// }

// func createReader(i interface{}, t *testing.T) io.Reader {
// 	data, err := json.Marshal(i)
// 	if err != nil {
// 		t.Fatal("cannot unmarshal:", err)
// 	}

// 	buf := bytes.Buffer{}
// 	_, err = buf.Write(data)
// 	if err != nil {
// 		t.Fatal("cannot write:", err)
// 	}

// 	return &buf
// }

// func TestUserHandler_Me(t *testing.T) {
// 	t.Skip("need bleve driver")
// 	router, handlers, f := createRouter(t)
// 	handler := handlers.UserHandler
// 	defer f()

// 	// Load fixtures
// 	user := &papernet.User{
// 		ID:        "1",
// 		Name:      "Test user",
// 		Bookmarks: []int{1},
// 	}
// 	err := handler.Repository.Upsert(user)
// 	if err != nil {
// 		t.Fatal("could not insert user:", err)
// 	}

// 	userToken, err := handler.Authenticator.Encoder.Encode(user.ID)
// 	if err != nil {
// 		t.Fatal("could not encode token:", err)
// 	}

// 	var tts = []struct {
// 		Token string
// 		Code  int
// 	}{
// 		{
// 			Token: "",
// 			Code:  401,
// 		},
// 		{
// 			Token: "not a bearer",
// 			Code:  401,
// 		},
// 		{
// 			Token: "bearer not.a.token",
// 			Code:  401,
// 		},
// 		{
// 			Token: fmt.Sprintf("bearer %s", userToken),
// 			Code:  200,
// 		},
// 	}

// 	for i, tt := range tts {
// 		req := httptest.NewRequest("GET", "/api/me", nil)
// 		req.Header.Add("authorization", tt.Token)
// 		resp := httptest.NewRecorder()
// 		router.ServeHTTP(resp, req)
// 		if resp.Code != tt.Code {
// 			t.Errorf("%d - incorrect code: expected %d got %d (body: %v)", i, tt.Code, resp.Code, resp.Body.String())
// 		}
// 	}
// }

// func TestUserHandler_UpdateBookmarks(t *testing.T) {
// 	t.Skip("need bleve driver")
// 	router, handlers, f := createRouter(t)
// 	handler := handlers.UserHandler
// 	defer f()

// 	// Load fixtures
// 	user := &papernet.User{
// 		ID:        "1",
// 		Name:      "Test user",
// 		Bookmarks: []int{1},
// 	}
// 	err := handler.Repository.Upsert(user)
// 	if err != nil {
// 		t.Fatal("could not insert user:", err)
// 	}

// 	userToken, err := handler.Authenticator.Encoder.Encode(user.ID)
// 	if err != nil {
// 		t.Fatal("could not encode token:", err)
// 	}

// 	var tts = []struct {
// 		Token    string
// 		Add      []int
// 		Remove   []int
// 		Code     int
// 		Expected []int
// 	}{
// 		{
// 			Token:    "",
// 			Add:      []int{},
// 			Remove:   []int{},
// 			Code:     401,
// 			Expected: []int{},
// 		},
// 		{
// 			Token:    fmt.Sprintf("bearer %s", userToken),
// 			Add:      []int{2},
// 			Remove:   []int{},
// 			Code:     200,
// 			Expected: []int{1, 2},
// 		},
// 		{
// 			Token:    fmt.Sprintf("bearer %s", userToken),
// 			Add:      []int{3},
// 			Remove:   []int{2},
// 			Code:     200,
// 			Expected: []int{1, 3},
// 		},
// 		{
// 			Token:    fmt.Sprintf("bearer %s", userToken),
// 			Add:      []int{2},
// 			Remove:   []int{2, 3},
// 			Code:     200,
// 			Expected: []int{1},
// 		},
// 	}

// 	for i, tt := range tts {
// 		reader := createReader(map[string]interface{}{
// 			"add":    tt.Add,
// 			"remove": tt.Remove,
// 		}, t)
// 		req := httptest.NewRequest("POST", "/api/bookmarks", reader)
// 		req.Header.Add("authorization", tt.Token)
// 		resp := httptest.NewRecorder()
// 		router.ServeHTTP(resp, req)
// 		if resp.Code != tt.Code {
// 			t.Errorf("%d - incorrect code: expected %d got %d (body: %v)", i, tt.Code, resp.Code, resp.Body.String())
// 			continue
// 		}

// 		if tt.Code >= 400 {
// 			continue
// 		}

// 		var r struct {
// 			Data papernet.User `json:"data"`
// 		}
// 		if err := json.Unmarshal(resp.Body.Bytes(), &r); err != nil {
// 			t.Errorf("%d - could not read json body: %v", i, err)
// 		} else if !reflect.DeepEqual(tt.Expected, r.Data.Bookmarks) {
// 			t.Errorf("%d - incorrect bookmarks: expected %v got %v", i, tt.Expected, r.Data.Bookmarks)
// 		}
// 	}
// }
<commit_msg>postpone testing :(<commit_after><|endoftext|>


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

// Package worker provides the workers that do the actual loading & processing
// of pages.
package worker

import (
	"fmt"
	"github.com/Matir/webborer/client"
	"github.com/Matir/webborer/logging"
	"github.com/Matir/webborer/results"
	ss "github.com/Matir/webborer/settings"
	"github.com/Matir/webborer/task"
	"github.com/Matir/webborer/util"
	"github.com/Matir/webborer/workqueue"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type Stoppable interface {
	Stop()
}

type PageWorker interface {
	Eligible(*http.Response) bool
	Handle(*task.Task, io.Reader)
}

// Workers do the work of connecting to the server, issuing the request, and
// then optionally parsing the response.  Normally a pool of several workers
// will be used due to network latency.
type Worker struct {
	// client for connections
	client client.Client
	// Channel for URLs to scan
	src <-chan *task.Task
	// Function to add future work
	adder workqueue.QueueAddFunc
	// Function to mark work done
	done workqueue.QueueDoneFunc
	// Channel for scan results
	rchan chan<- *results.Result
	// Settings
	settings *ss.ScanSettings
	// HTML worker to parse page
	pageWorker PageWorker
	// Channel to trigger stopping
	stop chan bool
	// Request for redirection
	redir *http.Request
	// Channel to signal worker stopping
	waitq chan bool
}

// Construct a worker with given settings.
func NewWorker(settings *ss.ScanSettings,
	factory client.ClientFactory,
	src <-chan *task.Task,
	adder workqueue.QueueAddFunc,
	done workqueue.QueueDoneFunc,
	rchan chan<- *results.Result) *Worker {
	w := &Worker{
		client:   factory.Get(),
		settings: settings,
		src:      src,
		adder:    adder,
		done:     done,
		rchan:    rchan,
		stop:     make(chan bool),
		waitq:    make(chan bool),
	}

	// Install redirect handler
	redirHandler := func(req *http.Request, _ []*http.Request) error {
		w.redir = req
		return fmt.Errorf("Stop redirect.")
	}
	w.client.SetCheckRedirect(redirHandler)

	return w
}

func (w *Worker) SetPageWorker(pw PageWorker) {
	w.pageWorker = pw
}

// Run the worker, processing input from a channel until either signalled to
// stop or the input channel is closed.
func (w *Worker) Run() {
	defer func() {
		w.waitq <- true
	}()
	for true {
		select {
		case <-w.stop:
			return
		case t, ok := <-w.src:
			if !ok { // channel closed
				return
			}
			w.HandleTask(t)
		}
	}
}

func (w *Worker) RunInBackground() {
	go w.Run()
}

func (w *Worker) Stop() {
	w.stop <- true
}

func (w *Worker) Wait() {
	<-w.waitq
}

func (w *Worker) HandleTask(t *task.Task) {
	logging.Logf(logging.LogDebug, "Trying Raw URL (unmangled): %s", t.String())
	withMangle := w.TryTask(t)
	if !util.URLIsDir(t.URL) {
		if withMangle {
			w.TryMangleTask(t)
		}
		if !util.URLHasExtension(t.URL) {
			for _, ext := range w.settings.Extensions {
				t := t.Copy()
				t.URL.Path += "." + ext
				if w.TryTask(t) {
					w.TryMangleTask(t)
				}
			}
		}
	}
	// Mark as done
	w.done(1)
}

func (w *Worker) TryMangleTask(t *task.Task) {
	if !w.settings.Mangle {
		return
	}
	clone := t.Copy()
	spos := strings.LastIndex(clone.URL.Path, "/")
	if spos == -1 {
		return
	}
	dirname := clone.URL.Path[:spos]
	basename := clone.URL.Path[spos+1:]
	for _, newname := range Mangle(basename) {
		clone := clone.Copy()
		clone.URL.Path = dirname + "/" + newname
		w.TryTask(clone)
	}
}

func (w *Worker) TryTask(t *task.Task) bool {
	logging.Logf(logging.LogInfo, "Trying: %s", t.String())
	tryMangle := false
	w.redir = nil
	if resp, err := w.client.Request(t.URL, t.Host, t.Header); err != nil && w.redir == nil {
		// TODO: add host, headers, group, etc.
		result := results.NewResult(t.URL, t.Host)
		result.Error = err
		if resp != nil {
			result.Code = resp.StatusCode
		}
		w.rchan <- result
	} else {
		defer resp.Body.Close()
		// Do we keep going?
		if util.URLIsDir(t.URL) && w.KeepSpidering(resp.StatusCode) {
			logging.Logf(logging.LogDebug, "Referring %s back for spidering.", t.String())
			w.adder(t)
		}
		if w.redir != nil {
			logging.Logf(logging.LogDebug, "Referring redirect %s back.", w.redir.URL.String())
			t := t.Copy()
			t.URL = w.redir.URL
			w.adder(t)
		}
		if w.pageWorker != nil && w.pageWorker.Eligible(resp) {
			w.pageWorker.Handle(t, resp.Body)
		}
		var redir *url.URL
		if w.redir != nil {
			redir = w.redir.URL
		}
		res := results.NewResult(t.URL, t.Host)
		res.Code = resp.StatusCode
		res.Redir = redir
		res.Length = resp.ContentLength
		res.ContentType = resp.Header.Get("Content-Type")
		w.rchan <- res
		tryMangle = w.KeepSpidering(resp.StatusCode)
	}
	if w.settings.SleepTime != 0 {
		time.Sleep(w.settings.SleepTime)
	}
	return tryMangle
}

// Should we keep spidering from this code?
func (w *Worker) KeepSpidering(code int) bool {
	if w.settings.RunMode == ss.RunModeDotProduct {
		return false
	}
	for _, v := range w.settings.SpiderCodes {
		if code == v {
			return true
		}
	}
	return false
}

// Starts a batch of workers based on the relevant settings.
func StartWorkers(settings *ss.ScanSettings,
	factory client.ClientFactory,
	src <-chan *task.Task,
	adder workqueue.QueueAddFunc,
	done workqueue.QueueDoneFunc,
	rchan chan<- *results.Result) []*Worker {
	count := settings.Workers
	workers := make([]*Worker, count)
	for i := 0; i < count; i++ {
		workers[i] = NewWorker(settings, factory, src, adder, done, rchan)
		workers[i].RunInBackground()
		if settings.ParseHTML {
			workers[i].SetPageWorker(NewHTMLWorker(adder))
		}
	}
	return workers
}

// Mangle a basename
func Mangle(basename string) []string {
	mangleRules := []string{
		".%s.swp", // VIM Swap File
		"%s~",     // Backup file
		"%s.bak",  // Backup file
		"%s.orig", // Backup file
	}
	res := make([]string, len(mangleRules))
	for i, rule := range mangleRules {
		res[i] = fmt.Sprintf(rule, basename)
	}
	return res
}
<commit_msg>Don't parse HTML except in enumeration mode.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package worker provides the workers that do the actual loading & processing
// of pages.
package worker

import (
	"fmt"
	"github.com/Matir/webborer/client"
	"github.com/Matir/webborer/logging"
	"github.com/Matir/webborer/results"
	ss "github.com/Matir/webborer/settings"
	"github.com/Matir/webborer/task"
	"github.com/Matir/webborer/util"
	"github.com/Matir/webborer/workqueue"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type Stoppable interface {
	Stop()
}

type PageWorker interface {
	Eligible(*http.Response) bool
	Handle(*task.Task, io.Reader)
}

// Workers do the work of connecting to the server, issuing the request, and
// then optionally parsing the response.  Normally a pool of several workers
// will be used due to network latency.
type Worker struct {
	// client for connections
	client client.Client
	// Channel for URLs to scan
	src <-chan *task.Task
	// Function to add future work
	adder workqueue.QueueAddFunc
	// Function to mark work done
	done workqueue.QueueDoneFunc
	// Channel for scan results
	rchan chan<- *results.Result
	// Settings
	settings *ss.ScanSettings
	// HTML worker to parse page
	pageWorker PageWorker
	// Channel to trigger stopping
	stop chan bool
	// Request for redirection
	redir *http.Request
	// Channel to signal worker stopping
	waitq chan bool
}

// Construct a worker with given settings.
func NewWorker(settings *ss.ScanSettings,
	factory client.ClientFactory,
	src <-chan *task.Task,
	adder workqueue.QueueAddFunc,
	done workqueue.QueueDoneFunc,
	rchan chan<- *results.Result) *Worker {
	w := &Worker{
		client:   factory.Get(),
		settings: settings,
		src:      src,
		adder:    adder,
		done:     done,
		rchan:    rchan,
		stop:     make(chan bool),
		waitq:    make(chan bool),
	}

	// Install redirect handler
	redirHandler := func(req *http.Request, _ []*http.Request) error {
		w.redir = req
		return fmt.Errorf("Stop redirect.")
	}
	w.client.SetCheckRedirect(redirHandler)

	return w
}

func (w *Worker) SetPageWorker(pw PageWorker) {
	w.pageWorker = pw
}

// Run the worker, processing input from a channel until either signalled to
// stop or the input channel is closed.
func (w *Worker) Run() {
	defer func() {
		w.waitq <- true
	}()
	for true {
		select {
		case <-w.stop:
			return
		case t, ok := <-w.src:
			if !ok { // channel closed
				return
			}
			w.HandleTask(t)
		}
	}
}

func (w *Worker) RunInBackground() {
	go w.Run()
}

func (w *Worker) Stop() {
	w.stop <- true
}

func (w *Worker) Wait() {
	<-w.waitq
}

func (w *Worker) HandleTask(t *task.Task) {
	logging.Logf(logging.LogDebug, "Trying Raw URL (unmangled): %s", t.String())
	withMangle := w.TryTask(t)
	if !util.URLIsDir(t.URL) {
		if withMangle {
			w.TryMangleTask(t)
		}
		if !util.URLHasExtension(t.URL) {
			for _, ext := range w.settings.Extensions {
				t := t.Copy()
				t.URL.Path += "." + ext
				if w.TryTask(t) {
					w.TryMangleTask(t)
				}
			}
		}
	}
	// Mark as done
	w.done(1)
}

func (w *Worker) TryMangleTask(t *task.Task) {
	if !w.settings.Mangle {
		return
	}
	clone := t.Copy()
	spos := strings.LastIndex(clone.URL.Path, "/")
	if spos == -1 {
		return
	}
	dirname := clone.URL.Path[:spos]
	basename := clone.URL.Path[spos+1:]
	for _, newname := range Mangle(basename) {
		clone := clone.Copy()
		clone.URL.Path = dirname + "/" + newname
		w.TryTask(clone)
	}
}

func (w *Worker) TryTask(t *task.Task) bool {
	logging.Logf(logging.LogInfo, "Trying: %s", t.String())
	tryMangle := false
	w.redir = nil
	if resp, err := w.client.Request(t.URL, t.Host, t.Header); err != nil && w.redir == nil {
		// TODO: add host, headers, group, etc.
		result := results.NewResult(t.URL, t.Host)
		result.Error = err
		if resp != nil {
			result.Code = resp.StatusCode
		}
		w.rchan <- result
	} else {
		defer resp.Body.Close()
		// Do we keep going?
		if util.URLIsDir(t.URL) && w.KeepSpidering(resp.StatusCode) {
			logging.Logf(logging.LogDebug, "Referring %s back for spidering.", t.String())
			w.adder(t)
		}
		if w.redir != nil {
			logging.Logf(logging.LogDebug, "Referring redirect %s back.", w.redir.URL.String())
			t := t.Copy()
			t.URL = w.redir.URL
			w.adder(t)
		}
		if w.pageWorker != nil && w.pageWorker.Eligible(resp) {
			w.pageWorker.Handle(t, resp.Body)
		}
		var redir *url.URL
		if w.redir != nil {
			redir = w.redir.URL
		}
		res := results.NewResult(t.URL, t.Host)
		res.Code = resp.StatusCode
		res.Redir = redir
		res.Length = resp.ContentLength
		res.ContentType = resp.Header.Get("Content-Type")
		w.rchan <- res
		tryMangle = w.KeepSpidering(resp.StatusCode)
	}
	if w.settings.SleepTime != 0 {
		time.Sleep(w.settings.SleepTime)
	}
	return tryMangle
}

// Should we keep spidering from this code?
func (w *Worker) KeepSpidering(code int) bool {
	if w.settings.RunMode == ss.RunModeDotProduct {
		return false
	}
	for _, v := range w.settings.SpiderCodes {
		if code == v {
			return true
		}
	}
	return false
}

// Starts a batch of workers based on the relevant settings.
func StartWorkers(settings *ss.ScanSettings,
	factory client.ClientFactory,
	src <-chan *task.Task,
	adder workqueue.QueueAddFunc,
	done workqueue.QueueDoneFunc,
	rchan chan<- *results.Result) []*Worker {
	count := settings.Workers
	workers := make([]*Worker, count)
	for i := 0; i < count; i++ {
		workers[i] = NewWorker(settings, factory, src, adder, done, rchan)
		workers[i].RunInBackground()
		if settings.ParseHTML && settings.RunMode == ss.RunModeEnumeration {
			workers[i].SetPageWorker(NewHTMLWorker(adder))
		}
	}
	return workers
}

// Mangle a basename
func Mangle(basename string) []string {
	mangleRules := []string{
		".%s.swp", // VIM Swap File
		"%s~",     // Backup file
		"%s.bak",  // Backup file
		"%s.orig", // Backup file
	}
	res := make([]string, len(mangleRules))
	for i, rule := range mangleRules {
		res[i] = fmt.Sprintf(rule, basename)
	}
	return res
}
<|endoftext|>


<commit_before>/*
Copyright 2016 caicloud authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package worker

import (
	"fmt"
	"strings"
	"time"

	"github.com/caicloud/cyclone/api"
	"github.com/caicloud/cyclone/cloud"
	"github.com/caicloud/cyclone/docker"
	"github.com/caicloud/cyclone/pkg/osutil"
	"github.com/caicloud/cyclone/websocket"
	"github.com/caicloud/cyclone/worker/ci"
	"github.com/caicloud/cyclone/worker/ci/parser"
	"github.com/caicloud/cyclone/worker/handler"
	"github.com/caicloud/cyclone/worker/helper"
	worker_log "github.com/caicloud/cyclone/worker/log"
	"github.com/caicloud/cyclone/worker/vcs"
	"github.com/zoumo/logdog"
)

const (
	WorkerTimeout = 7200 * time.Second
	WaitTimes     = 5

	DefaultCaicloudYaml = "build:\n  dockerfile_name: Dockerfile"
)

// Worker ...
type Worker struct {
	Config *Config
	Envs   *cloud.WorkerEnvs
}

// Run starts the worker
func (worker *Worker) Run() error {

	logdog.Info("worker start with", logdog.Fields{"config": worker.Config, "envs": worker.Envs})

	// Get event for cyclone server
	event, err := worker.getEvent()
	if err != nil {
		logdog.Errorf("get event err: %v", err)
		event.Status = api.EventStatusFail
		event.ErrorMessage = err.Error()
		sendErr := worker.sendEvent(event)
		if sendErr != nil {
			logdog.Errorf("set event result err: %v", err)
		}
		return err
	}
	logdog.Info("get event success", logdog.Fields{"event": event})

	// Handle event
	logdog.Info("handleEvent ...")
	worker.handleEvent(&event)

	// Sent event for cyclone server
	err = worker.sendEvent(event)
	if err != nil {
		logdog.Errorf("set event result err: %v", err)
		return err
	}
	logdog.Info("send event to server", logdog.Fields{"event": event})
	return nil
}

// getEvent used for getting event for cyclone server
func (worker *Worker) getEvent() (api.Event, error) {
	eventID := worker.Config.ID
	serverHost := worker.Envs.CycloneServer

	BaseURL := fmt.Sprintf("%s/api/%s", serverHost, api.APIVersion)
	httpHandler := handler.NewHTTPHandler(BaseURL)

	var getResponse api.GetEventResponse
	err := httpHandler.GetEvent(eventID, &getResponse)
	if err != nil {
		return api.Event{}, err
	}

	event := getResponse.Event

	return event, nil
}

// handleEvent analysize the the operation in event, and do the relate operation
func (worker *Worker) handleEvent(event *api.Event) {
	vcsManager := vcs.NewManager()

	logServer := worker.Envs.LogServer
	err := worker_log.DialLogServer(logServer)
	if err != nil {
		logdog.Errorf("dail log server err: %v", err)
	} else {
		go worker_log.SendHeartBeat()
		defer worker_log.Disconnect()
	}

	switch event.Operation {
	case "create-service":
		worker.createService(vcsManager, event)

	case "create-version":
		worker.createVersion(vcsManager, event)

	default:
		event.Status = api.EventStatusFail
		event.ErrorMessage = "unkwon operation"
		logdog.Error("Operation failed", logdog.Fields{"event": event})
	}
}

// createService verify repo validity for creating service
func (worker *Worker) createService(vcsManager *vcs.Manager, event *api.Event) {
	// err := vcsManager.CloneServiceRepository(event)
	err := vcsManager.CheckRepoValid(event)
	if err != nil {
		event.Status = api.EventStatusFail
		event.ErrorMessage = err.Error()
		logdog.Error("Operation failed", logdog.Fields{"event": event})
	} else {
		event.Status = api.EventStatusSuccess
	}
}

// createVersion create version for service and push log to server via websocket
// Step1: clone code from VCS
// Step2: integretion
// Step3: publish
// Step4: deploy
func (worker *Worker) createVersion(vcsManager *vcs.Manager, event *api.Event) {
	registryLocation := worker.Envs.RegistryLocation
	registryUsername := worker.Envs.RegistryUsername
	registryPassword := worker.Envs.RegistryPassword

	dockerManager, err := docker.NewManager(worker.Config.DockerHost, "",
		api.RegistryCompose{registryLocation, registryUsername, registryPassword})
	if err != nil {
		return
	}
	ciManager, err := ci.NewManager(dockerManager)
	if err != nil {
		return
	}

	err = worker_log.CreateFileBuffer(event.EventID)
	if err != nil {
		event.Status = api.EventStatusFail
		event.ErrorMessage = err.Error()
		logdog.Error("Operation failed", logdog.Fields{"event": event})
		return
	}
	output := worker_log.Output
	ch := make(chan interface{})
	defer func() {
		// Push log file to cyclone.
		if err := helper.PushLogToCyclone(event); err != nil {
			event.Status = api.EventStatusFail
			event.ErrorMessage = err.Error()
			logdog.Error("Operation failed", logdog.Fields{"event": event})
		}
		output.Close()
		worker_log.SetWatchLogFileSwitch(output.Name(), false)
		// wait util the send the log to kafka throuth cyclone server totally.
		for i := 0; i < WaitTimes; i++ {
			if isChanClosed(ch) {
				break
			}
			time.Sleep(time.Second * 2)
		}
	}()

	topicLog := websocket.CreateTopicName(string(event.Operation), event.Service.UserID,
		event.Service.ServiceID, event.Version.VersionID)
	go worker_log.WatchLogFile(output.Name(), topicLog, ch)

	destDir := vcsManager.GetCloneDir(&event.Service, &event.Version)
	event.Data["context-dir"] = destDir
	event.Data["image-name"] = fmt.Sprintf("%s/%s/%s", dockerManager.Registry,
		strings.ToLower(event.Service.Username), strings.ToLower(event.Service.Name))
	event.Data["tag-name"] = event.Version.Name

	if err = vcsManager.CloneVersionRepository(event); err != nil {
		event.Status = api.EventStatusFail
		event.ErrorMessage = err.Error()
		logdog.Error("Operation failed", logdog.Fields{"event": event})
		return
	}

	replaceDockerfile(event, destDir)
	replaceCaicloudYaml(event, destDir)

	// Get the execution tree from the caicloud.yml.
	tree, err := ciManager.Parse(event)
	if err != nil {
		event.Status = api.EventStatusFail
		event.ErrorMessage = err.Error()
		logdog.Error("Operation failed", logdog.Fields{"event": event})
		return

	}
	worker.yamlBuild(event, tree, dockerManager, ciManager)
}

// replaceDockerfile create new dockerfile to replace the dockerfile in repo
func replaceDockerfile(event *api.Event, destDir string) {

	// Create dockerfile and caicloud.yml if need
	if event.Service.Dockerfile != "" {
		path := destDir + "/Dockerfile"
		err := osutil.ReplaceFile(path, strings.NewReader(event.Service.Dockerfile))
		if err != nil {
			worker_log.InsertStepLog(event, worker_log.CloneRepository, worker_log.Stop, err)
			return
		}
	}
}

// replaceCaicloudYaml create new yaml to replace the yaml in repo
// default yaml file is caicloud.yml
// if you set your own yaml config name, cyclone will use it.
// if there is no yaml found, cyclone will create a default yaml file with only build step.
func replaceCaicloudYaml(event *api.Event, destDir string) {

	yamlFile := destDir + "/" + ci.DefaultYamlFile
	if event.Service.YAMLConfigName != "" {
		yamlFile = destDir + "/" + event.Service.YAMLConfigName
	}

	if event.Service.CaicloudYaml != "" || !osutil.IsFileExists(yamlFile) {
		content := DefaultCaicloudYaml
		if event.Service.CaicloudYaml != "" {
			content = event.Service.CaicloudYaml
		}
		err := osutil.ReplaceFile(yamlFile, strings.NewReader(content))
		if err != nil {
			worker_log.InsertStepLog(event, worker_log.CloneRepository, worker_log.Stop, err)
			return
		}
	}
}

// yamlBuild func uses for build with yaml file.
func (worker *Worker) yamlBuild(event *api.Event, tree *parser.Tree, dockerManager *docker.Manager, ciManager *ci.Manager) {
	operation := string(event.Version.Operation)
	// Load the tree to a runner.Build.
	r, err := ciManager.LoadTree(event, tree)
	if err != nil {
		event.Status = api.EventStatusFail
		event.ErrorMessage = err.Error()
		logdog.Error("Operation failed", logdog.Fields{"event": event})
		return
	}

	// Build image
	if err = helper.ExecBuild(ciManager, r); err != nil {
		event.Status = api.EventStatusFail
		event.ErrorMessage = err.Error()
		logdog.Error("Operation failed", logdog.Fields{"event": event})
		return
	}

	// If need integration
	if strings.Contains(operation, "integration") {
		// Integration
		if err = helper.ExecIntegration(ciManager, r); err != nil {
			event.Status = api.EventStatusFail
			event.ErrorMessage = err.Error()
			logdog.Error("Operation failed", logdog.Fields{"event": event})
			return
		}
	}

	// If need publish
	if strings.Contains(operation, "publish") {
		// Publish
		if err = helper.ExecPublish(ciManager, r); err != nil {
			event.Status = api.EventStatusFail
			event.ErrorMessage = err.Error()
			logdog.Error("Operation failed", logdog.Fields{"event": event})
			return
		}
	}

	// If need deploy
	if strings.Contains(operation, "deploy") {
		// Deploy
		if err = helper.ExecDeploy(event, dockerManager, r, tree); err != nil {
			event.Status = api.EventStatusFail
			event.ErrorMessage = err.Error()
			logdog.Error("Operation failed", logdog.Fields{"event": event})
			return
		}

		err = worker.sendEvent(*event)
		if err != nil {
			logdog.Errorf("set event result err: %v", err)
		}

		// Deploy Check
		helper.ExecDeployCheck(event, tree)

	}
	event.Status = api.EventStatusSuccess
}

// sendEvent used for setting event for circe server
func (worker *Worker) sendEvent(event api.Event) error {
	eventID := worker.Config.ID
	serverHost := worker.Envs.CycloneServer

	BaseURL := fmt.Sprintf("%s/api/%s", serverHost, api.APIVersion)
	httpHandler := handler.NewHTTPHandler(BaseURL)
	result := &api.SetEvent{
		Event: event,
	}
	var setResponse api.SetEventResponse
	var err error
	DueTime := time.Now().Add(time.Duration(WorkerTimeout))
	for DueTime.After(time.Now()) == true {
		err = httpHandler.SetEvent(eventID, result, &setResponse)
		if err != nil {
			if strings.Contains(err.Error(), "connection refused") {
				time.Sleep(time.Minute)
				continue
			}
			return err
		}
		logdog.Infof("set event result: %v", setResponse)
		return nil
	}
	return err
}

func isChanClosed(ch chan interface{}) bool {
	if len(ch) == 0 {
		select {
		case _, ok := <-ch:
			return !ok
		}
	}
	return false
}
<commit_msg>refactor(woker): add function to set fail status for event (#168)<commit_after>/*
Copyright 2016 caicloud authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package worker

import (
	"fmt"
	"strings"
	"time"

	"github.com/caicloud/cyclone/api"
	"github.com/caicloud/cyclone/cloud"
	"github.com/caicloud/cyclone/docker"
	"github.com/caicloud/cyclone/pkg/osutil"
	"github.com/caicloud/cyclone/websocket"
	"github.com/caicloud/cyclone/worker/ci"
	"github.com/caicloud/cyclone/worker/ci/parser"
	"github.com/caicloud/cyclone/worker/handler"
	"github.com/caicloud/cyclone/worker/helper"
	worker_log "github.com/caicloud/cyclone/worker/log"
	"github.com/caicloud/cyclone/worker/vcs"
	"github.com/zoumo/logdog"
)

const (
	WorkerTimeout = 7200 * time.Second
	WaitTimes     = 5

	DefaultCaicloudYaml = "build:\n  dockerfile_name: Dockerfile"
)

// Worker ...
type Worker struct {
	Config *Config
	Envs   *cloud.WorkerEnvs
}

// Run starts the worker
func (worker *Worker) Run() error {

	logdog.Info("worker start with", logdog.Fields{"config": worker.Config, "envs": worker.Envs})

	// Get event for cyclone server
	event, err := worker.getEvent()
	if err != nil {
		setEventFailStatus(&event, err.Error())
		sendErr := worker.sendEvent(event)
		if sendErr != nil {
			logdog.Errorf("set event result err: %v", err)
		}
		return err
	}
	logdog.Info("get event success", logdog.Fields{"event": event})

	// Handle event
	logdog.Info("handleEvent ...")
	worker.handleEvent(&event)

	// Sent event for cyclone server
	err = worker.sendEvent(event)
	if err != nil {
		logdog.Errorf("set event result err: %v", err)
		return err
	}
	logdog.Info("send event to server", logdog.Fields{"event": event})
	return nil
}

// getEvent used for getting event for cyclone server
func (worker *Worker) getEvent() (api.Event, error) {
	eventID := worker.Config.ID
	serverHost := worker.Envs.CycloneServer

	BaseURL := fmt.Sprintf("%s/api/%s", serverHost, api.APIVersion)
	httpHandler := handler.NewHTTPHandler(BaseURL)

	var getResponse api.GetEventResponse
	err := httpHandler.GetEvent(eventID, &getResponse)
	if err != nil {
		return api.Event{}, err
	}

	event := getResponse.Event

	return event, nil
}

// handleEvent analysize the the operation in event, and do the relate operation
func (worker *Worker) handleEvent(event *api.Event) {
	vcsManager := vcs.NewManager()

	logServer := worker.Envs.LogServer
	err := worker_log.DialLogServer(logServer)
	if err != nil {
		logdog.Errorf("dail log server err: %v", err)
	} else {
		go worker_log.SendHeartBeat()
		defer worker_log.Disconnect()
	}

	switch event.Operation {
	case "create-service":
		worker.createService(vcsManager, event)

	case "create-version":
		worker.createVersion(vcsManager, event)

	default:
		setEventFailStatus(event, "unkwon operation")
	}
}

// createService verify repo validity for creating service
func (worker *Worker) createService(vcsManager *vcs.Manager, event *api.Event) {
	// err := vcsManager.CloneServiceRepository(event)
	err := vcsManager.CheckRepoValid(event)
	if err != nil {
		setEventFailStatus(event, err.Error())
	} else {
		event.Status = api.EventStatusSuccess
	}

}

// createVersion create version for service and push log to server via websocket
// Step1: clone code from VCS
// Step2: integretion
// Step3: publish
// Step4: deploy
func (worker *Worker) createVersion(vcsManager *vcs.Manager, event *api.Event) {
	registryLocation := worker.Envs.RegistryLocation
	registryUsername := worker.Envs.RegistryUsername
	registryPassword := worker.Envs.RegistryPassword

	dockerManager, err := docker.NewManager(worker.Config.DockerHost, "",
		api.RegistryCompose{registryLocation, registryUsername, registryPassword})
	if err != nil {
		setEventFailStatus(event, err.Error())
		return
	}
	ciManager, err := ci.NewManager(dockerManager)
	if err != nil {
		setEventFailStatus(event, err.Error())
		return
	}

	err = worker_log.CreateFileBuffer(event.EventID)
	if err != nil {
		setEventFailStatus(event, err.Error())
		return
	}
	output := worker_log.Output
	ch := make(chan interface{})
	defer func() {
		// Push log file to cyclone.
		if err := helper.PushLogToCyclone(event); err != nil {
			setEventFailStatus(event, err.Error())
		}
		output.Close()
		worker_log.SetWatchLogFileSwitch(output.Name(), false)
		// wait util the send the log to kafka throuth cyclone server totally.
		for i := 0; i < WaitTimes; i++ {
			if isChanClosed(ch) {
				break
			}
			time.Sleep(time.Second * 2)
		}
	}()

	topicLog := websocket.CreateTopicName(string(event.Operation), event.Service.UserID,
		event.Service.ServiceID, event.Version.VersionID)
	go worker_log.WatchLogFile(output.Name(), topicLog, ch)

	destDir := vcsManager.GetCloneDir(&event.Service, &event.Version)
	event.Data["context-dir"] = destDir
	event.Data["image-name"] = fmt.Sprintf("%s/%s/%s", dockerManager.Registry,
		strings.ToLower(event.Service.Username), strings.ToLower(event.Service.Name))
	event.Data["tag-name"] = event.Version.Name

	if err = vcsManager.CloneVersionRepository(event); err != nil {
		setEventFailStatus(event, err.Error())
		return
	}

	replaceDockerfile(event, destDir)
	replaceCaicloudYaml(event, destDir)

	// Get the execution tree from the caicloud.yml.
	tree, err := ciManager.Parse(event)
	if err != nil {
		setEventFailStatus(event, err.Error())
		return

	}
	worker.yamlBuild(event, tree, dockerManager, ciManager)
}

// replaceDockerfile create new dockerfile to replace the dockerfile in repo
func replaceDockerfile(event *api.Event, destDir string) {

	// Create dockerfile and caicloud.yml if need
	if event.Service.Dockerfile != "" {
		path := destDir + "/Dockerfile"
		err := osutil.ReplaceFile(path, strings.NewReader(event.Service.Dockerfile))
		if err != nil {
			worker_log.InsertStepLog(event, worker_log.CloneRepository, worker_log.Stop, err)
			return
		}
	}
}

// replaceCaicloudYaml create new yaml to replace the yaml in repo
// default yaml file is caicloud.yml
// if you set your own yaml config name, cyclone will use it.
// if there is no yaml found, cyclone will create a default yaml file with only build step.
func replaceCaicloudYaml(event *api.Event, destDir string) {

	yamlFile := destDir + "/" + ci.DefaultYamlFile
	if event.Service.YAMLConfigName != "" {
		yamlFile = destDir + "/" + event.Service.YAMLConfigName
	}

	if event.Service.CaicloudYaml != "" || !osutil.IsFileExists(yamlFile) {
		content := DefaultCaicloudYaml
		if event.Service.CaicloudYaml != "" {
			content = event.Service.CaicloudYaml
		}
		err := osutil.ReplaceFile(yamlFile, strings.NewReader(content))
		if err != nil {
			worker_log.InsertStepLog(event, worker_log.CloneRepository, worker_log.Stop, err)
			return
		}
	}
}

// yamlBuild func uses for build with yaml file.
func (worker *Worker) yamlBuild(event *api.Event, tree *parser.Tree, dockerManager *docker.Manager, ciManager *ci.Manager) {
	operation := string(event.Version.Operation)
	// Load the tree to a runner.Build.
	r, err := ciManager.LoadTree(event, tree)
	if err != nil {
		setEventFailStatus(event, err.Error())
		return
	}

	// Build image
	if err = helper.ExecBuild(ciManager, r); err != nil {
		setEventFailStatus(event, err.Error())
		return
	}

	// If need integration
	if strings.Contains(operation, "integration") {
		// Integration
		if err = helper.ExecIntegration(ciManager, r); err != nil {
			setEventFailStatus(event, err.Error())
			return
		}
	}

	// If need publish
	if strings.Contains(operation, "publish") {
		// Publish
		if err = helper.ExecPublish(ciManager, r); err != nil {
			setEventFailStatus(event, err.Error())
			return
		}
	}

	// If need deploy
	if strings.Contains(operation, "deploy") {
		// Deploy
		if err = helper.ExecDeploy(event, dockerManager, r, tree); err != nil {
			setEventFailStatus(event, err.Error())
			return
		}

		err = worker.sendEvent(*event)
		if err != nil {
			logdog.Errorf("set event result err: %v", err)
		}

		// Deploy Check
		helper.ExecDeployCheck(event, tree)

	}
	event.Status = api.EventStatusSuccess
}

// sendEvent used for setting event for circe server
func (worker *Worker) sendEvent(event api.Event) error {
	eventID := worker.Config.ID
	serverHost := worker.Envs.CycloneServer

	BaseURL := fmt.Sprintf("%s/api/%s", serverHost, api.APIVersion)
	httpHandler := handler.NewHTTPHandler(BaseURL)
	result := &api.SetEvent{
		Event: event,
	}
	var setResponse api.SetEventResponse
	var err error
	DueTime := time.Now().Add(time.Duration(WorkerTimeout))
	for DueTime.After(time.Now()) == true {
		err = httpHandler.SetEvent(eventID, result, &setResponse)
		if err != nil {
			if strings.Contains(err.Error(), "connection refused") {
				time.Sleep(time.Minute)
				continue
			}
			return err
		}
		logdog.Infof("set event result: %v", setResponse)
		return nil
	}
	return err
}

func isChanClosed(ch chan interface{}) bool {
	if len(ch) == 0 {
		select {
		case _, ok := <-ch:
			return !ok
		}
	}
	return false
}

// setEventFailStatus sets the fail status of the event.
func setEventFailStatus(event *api.Event, ErrorMessage string) {
	event.Status = api.EventStatusFail
	event.ErrorMessage = ErrorMessage
	logdog.Error("Operation failed", logdog.Fields{"event": event})
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/bmizerany/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
	Notes                                                                                []map[string]interface{}
}

type Worker struct {
	Db        *sql.DB
	LastRunAt time.Time
}

var worker Worker
var sr_number string

func init() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	worker.Db = db

	// fetch SR num from command line, if present
	flag.StringVar(&sr_number, "sr-number", "", "SR number to fetch")
}

func main() {
	defer worker.Db.Close()
	flag.Parse()

        if sr_number != "" {
                sr := fetchSingleRequest(sr_number)
		sr.Save()
		return
        }

	for {
		switch {
		case time.Since(worker.LastRunAt) > (30 * time.Second):
			// load requests from open311
			for _, request := range fetchRequests() {
				request.Save()
			}
			worker.LastRunAt = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func (req Open311Request) Save() (persisted bool) {
	// create or update a SR

	// open311 says we should always ignore a SR that does not have a SR# assigned
	if req.Service_request_id == "" {
		log.Printf("cowardly refusing to create a new SR record because of empty SR#. Request type is %s", req.Service_name)
		return false
	}

	persisted = false

	// find existing record if exists
	var existing_id int
	err := worker.Db.QueryRow("SELECT id FROM service_requests WHERE service_request_id = $1", req.Service_request_id).Scan(&existing_id)
	switch {
	case err == sql.ErrNoRows:
		// log.Printf("did not find existing record %s", req.Service_request_id)
	case err != nil:
		// log.Print("error searching for existing SR", err)
	default:
		persisted = true
		// log.Printf("found existing sr %s", req.Service_request_id)
	}

	var stmt *sql.Stmt

	if !persisted {
		// create new record
		stmt, err = worker.Db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id, closed_datetime) " +
			"VALUES ($1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17); ")

		// "WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

	} else {
		// update existing record
		stmt, err = worker.Db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW(), closed_datetime = $17 WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}
	}

	tx, err := worker.Db.Begin()

	if err != nil {
		log.Fatal("error beginning transaction", err)
	}

	var closed_time pq.NullTime
	if t := req.ExtractClosedDatetime(); !t.IsZero() {
		closed_time.Time = t
	}

	log.Printf("closed time for SR %s is %+v", req.Service_request_id, closed_time)

	_, err = tx.Stmt(stmt).Exec(req.Service_request_id,
		req.Status,
		req.Service_name,
		req.Service_code,
		req.Agency_responsible,
		req.Address,
		req.Requested_datetime,
		req.Updated_datetime,
		req.Lat,
		req.Long,
		req.Extended_attributes["ward"],
		req.Extended_attributes["police_district"],
		req.Media_url,
		req.Extended_attributes["channel"],
		req.Extended_attributes["duplicate"],
		req.Extended_attributes["parent_service_request_id"],
		closed_time,
	)

	if err != nil {
		log.Fatalf("could not update %s because %s", req.Service_request_id, err)
	} else {
		var verb string
		switch {
		case !persisted && closed_time.Time.IsZero():
			verb = "CREATED"
		case !persisted && !closed_time.Time.IsZero():
			verb = "CREATED/CLOSED"
		case persisted && closed_time.Time.IsZero():
			verb = "UPDATED"
		case persisted && !closed_time.Time.IsZero():
			verb = "UPDATED/CLOSED"
		}

		log.Printf("[%s] %s", verb, req)
		persisted = true
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal("error closing transaction", err)
	}

	return persisted

	// calculate closed time if necessary

}

func (req Open311Request) ExtractClosedDatetime() time.Time {
	// given an extended_attributes JSON blob, pluck out the closed time, if present
	// req.PrintNotes()

	var closed_at time.Time
	for _, note := range req.Notes {
		if note["type"] == "closed" {
			parsed_date, err := time.Parse("2006-01-02T15:04:05-07:00", note["datetime"].(string))
			if err != nil {
				log.Print("error parsing date", err)
			}
			log.Printf("located closed date %s for sr: %s", parsed_date, req)
			closed_at = parsed_date
		}
	}

	return closed_at
}

func (req Open311Request) PrintNotes() {
	fmt.Printf("Notes for SR %s:\n", req.Service_request_id)

	for _, note := range req.Notes {
		fmt.Printf("%+v\n", note)
	}
}

func fetchSingleRequest(sr_number string) (request Open311Request) {
	// given an SR, fetch the record
	log.Printf("fetching single SR %s", sr_number)
	open311_api_endpoint := fmt.Sprintf("http://311api.cityofchicago.org/open311/v2/requests/%s.json?extensions=true", sr_number)

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	var requests []Open311Request

	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests[0]
}

func fetchRequests() (requests []Open311Request) {
	// find the most recent SR that we know about in the database
	rows, err := worker.Db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Print("error finding most recent SR", err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}
<commit_msg>fix bug that prevented closed_datetime from being set on closed SR gofmt tweaks<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/bmizerany/pq"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const OPEN311_API_URI = "http://311api.cityofchicago.org/open311/v2/requests.json?extensions=true&page_size=500"

type Open311Request struct {
	Lat, Long                                                                                float64
	Ward, Police_district                                                                                int
	Service_request_id, Status, Service_name, Service_code, Agency_responsible, Address, Channel, Media_url string
	Requested_datetime, Updated_datetime                                                                    string // FIXME: should these be proper time objects?
	Extended_attributes                                                                                map[string]interface{}
	Notes                                                                                []map[string]interface{}
}

type Worker struct {
	Db        *sql.DB
	LastRunAt time.Time
}

var worker Worker
var sr_number string

func init() {
	// open database
	db, err := sql.Open("postgres", "dbname=cwfy sslmode=disable")
	if err != nil {
		log.Fatal("Cannot open database connection", err)
	}
	worker.Db = db

	// fetch SR num from command line, if present
	flag.StringVar(&sr_number, "sr-number", "", "SR number to fetch")
}

func main() {
	defer worker.Db.Close()
	flag.Parse()

	if sr_number != "" {
		sr := fetchSingleRequest(sr_number)
		sr.Save()
		return
	}

	for {
		switch {
		case time.Since(worker.LastRunAt) > (30 * time.Second):
			// load requests from open311
			for _, request := range fetchRequests() {
				request.Save()
			}
			worker.LastRunAt = time.Now()
		default:
			log.Print("sleeping for 10 seconds")
			time.Sleep(10 * time.Second)
		}
	}
}

func (req Open311Request) String() string {
	// pretty print SR information
	return fmt.Sprintf("%s: %s at %s %f,%f, last update %s", req.Service_request_id, req.Service_name, req.Address, req.Lat, req.Long, req.Updated_datetime)
}

func (req Open311Request) Save() (persisted bool) {
	// create or update a SR

	// open311 says we should always ignore a SR that does not have a SR# assigned
	if req.Service_request_id == "" {
		log.Printf("cowardly refusing to create a new SR record because of empty SR#. Request type is %s", req.Service_name)
		return false
	}

	persisted = false

	// find existing record if exists
	var existing_id int
	err := worker.Db.QueryRow("SELECT id FROM service_requests WHERE service_request_id = $1", req.Service_request_id).Scan(&existing_id)
	switch {
	case err == sql.ErrNoRows:
		// log.Printf("did not find existing record %s", req.Service_request_id)
	case err != nil:
		// log.Print("error searching for existing SR", err)
	default:
		persisted = true
		// log.Printf("found existing sr %s", req.Service_request_id)
	}

	var stmt *sql.Stmt

	if !persisted {
		// create new record
		stmt, err = worker.Db.Prepare("INSERT INTO service_requests(service_request_id," +
			"status, service_name, service_code, agency_responsible, " +
			"address, requested_datetime, updated_datetime, lat, long," +
			"ward, police_district, media_url, channel, duplicate, parent_service_request_id, closed_datetime) " +
			"VALUES ($1::varchar, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17); ")

		// "WHERE NOT EXISTS (SELECT 1 FROM service_requests WHERE service_request_id = $1);")

		if err != nil {
			log.Fatal("error preparing database insert statement", err)
		}

	} else {
		// update existing record
		stmt, err = worker.Db.Prepare("UPDATE service_requests SET " +
			"status = $2, service_name = $3, service_code = $4, agency_responsible = $5, " +
			"address = $6, requested_datetime = $7, updated_datetime = $8, lat = $9, long = $10," +
			"ward = $11, police_district = $12, media_url = $13, channel = $14, duplicate = $15, " +
			"parent_service_request_id = $16, updated_at = NOW(), closed_datetime = $17 WHERE service_request_id = $1;")

		if err != nil {
			log.Fatal("error preparing database update statement", err)
		}
	}

	tx, err := worker.Db.Begin()

	if err != nil {
		log.Fatal("error beginning transaction", err)
	}

	t := req.ExtractClosedDatetime()
	closed_time := pq.NullTime{Time: t, Valid: !t.IsZero()}

	_, err = tx.Stmt(stmt).Exec(req.Service_request_id,
		req.Status,
		req.Service_name,
		req.Service_code,
		req.Agency_responsible,
		req.Address,
		req.Requested_datetime,
		req.Updated_datetime,
		req.Lat,
		req.Long,
		req.Extended_attributes["ward"],
		req.Extended_attributes["police_district"],
		req.Media_url,
		req.Extended_attributes["channel"],
		req.Extended_attributes["duplicate"],
		req.Extended_attributes["parent_service_request_id"],
		closed_time,
	)

	if err != nil {
		log.Fatalf("could not update %s because %s", req.Service_request_id, err)
	} else {
		var verb string
		switch {
		case !persisted && closed_time.Time.IsZero():
			verb = "CREATED"
		case !persisted && !closed_time.Time.IsZero():
			verb = "CREATED/CLOSED"
		case persisted && closed_time.Time.IsZero():
			verb = "UPDATED"
		case persisted && !closed_time.Time.IsZero():
			verb = "UPDATED/CLOSED"
		}

		log.Printf("[%s] %s", verb, req)
		persisted = true
	}

	err = tx.Commit()
	if err != nil {
		log.Fatal("error closing transaction", err)
	}

	return persisted

	// calculate closed time if necessary

}

func (req Open311Request) ExtractClosedDatetime() time.Time {
	// given an extended_attributes JSON blob, pluck out the closed time, if present
	// req.PrintNotes()

	var closed_at time.Time
	for _, note := range req.Notes {
		if note["type"] == "closed" {
			parsed_date, err := time.Parse("2006-01-02T15:04:05-07:00", note["datetime"].(string))
			if err != nil {
				log.Print("error parsing date", err)
			}
			log.Printf("SR %s closed at: %s", req, parsed_date)
			closed_at = parsed_date
		}
	}

	return closed_at
}

func (req Open311Request) PrintNotes() {
	fmt.Printf("Notes for SR %s:\n", req.Service_request_id)

	for _, note := range req.Notes {
		fmt.Printf("%+v\n", note)
	}
}

func fetchSingleRequest(sr_number string) (request Open311Request) {
	// given an SR, fetch the record
	log.Printf("fetching single SR %s", sr_number)
	open311_api_endpoint := fmt.Sprintf("http://311api.cityofchicago.org/open311/v2/requests/%s.json?extensions=true", sr_number)

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	var requests []Open311Request

	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests[0]
}

func fetchRequests() (requests []Open311Request) {
	// find the most recent SR that we know about in the database
	rows, err := worker.Db.Query("SELECT MAX(updated_datetime) FROM service_requests;")
	if err != nil {
		log.Fatal("error finding most recent service request", err)
	}

	last_updated_at := time.Now()
	for rows.Next() {
		if err := rows.Scan(&last_updated_at); err != nil {
			log.Print("error finding most recent SR", err)
		}

		log.Printf("most recent SR timestamp %s", last_updated_at)
	}

	// janky hack to transform the last updated timestamp into
	// a format that plays nicely with the Open311 API
	// FIXME: there HAS to be a better way to handle this.
	formatted_date_string := last_updated_at.Format(time.RFC3339)
	formatted_date_string_with_tz := formatted_date_string[0:len(formatted_date_string)-1] + "-0500" // trunc the trailing 'Z' and tack on timezone

	// construct the request URI using base params and the proper time
	open311_api_endpoint := OPEN311_API_URI + "&updated_after=" + formatted_date_string_with_tz

	log.Printf("fetching from %s", open311_api_endpoint)
	resp, err := http.Get(open311_api_endpoint)
	defer resp.Body.Close()

	if err != nil {
		log.Fatalln("error fetching from Open311 endpoint", err)
	}

	// load response body
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("error loading response body", err)
	}

	// parse JSON and load into an array of Open311Request objects
	err = json.Unmarshal(body, &requests)
	if err != nil {
		log.Fatal("error parsing JSON:", err)
	}

	log.Printf("received %d requests from Open311", len(requests))

	return requests
}
<|endoftext|>


<commit_before>package uploader

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"time"
)

type Uploader interface {
	Upload(u *url.URL, filename string, r *http.Request) (*http.Response, *url.URL, error)
	Poll(u *url.URL, res *http.Response, closeChan <-chan bool, interval time.Duration) error
}

type httpUploader struct {
	baseUrl *url.URL
	client  *http.Client
}

func New(baseUrl *url.URL, transport http.RoundTripper) Uploader {
	return &httpUploader{
		baseUrl: baseUrl,
		client: &http.Client{
			Transport: transport,
		},
	}
}

var headersToForward = []string{
	"Content-MD5",
}

func (u *httpUploader) Upload(primaryUrl *url.URL, filename string, r *http.Request) (*http.Response, *url.URL, error) {
	if r.ContentLength <= 0 {
		return &http.Response{StatusCode: http.StatusLengthRequired}, nil, fmt.Errorf("Missing Content Length")
	}
	defer r.Body.Close()

	uploadReq, err := newMultipartRequestFromReader(r.ContentLength, r.Body, "upload[droplet]", filename)
	if err != nil {
		return nil, nil, err
	}
	for _, headerName := range headersToForward {
		uploadReq.Header.Set(headerName, r.Header.Get(headerName))
	}

	// try the fast path
	uploadReq.URL = primaryUrl
	if primaryUrl.User != nil {
		if password, set := primaryUrl.User.Password(); set {
			uploadReq.SetBasicAuth(primaryUrl.User.Username(), password)
		}
	}

	rsp, err := u.do(uploadReq)
	if err == nil {
		return rsp, uploadReq.URL, err
	}

	// not a connect (dial) error
	var nestedErr error = err
	if urlErr, ok := err.(*url.Error); ok {
		nestedErr = urlErr.Err
	}
	if netErr, ok := nestedErr.(*net.OpError); !ok || netErr.Op != "dial" {
		return rsp, nil, err
	}

	// try the slow path
	uploadReq, err = newMultipartRequestFromReader(r.ContentLength, r.Body, "upload[droplet]", filename)
	if err != nil {
		return nil, nil, err
	}

	uploadReq.URL.Scheme = u.baseUrl.Scheme
	uploadReq.URL.Host = u.baseUrl.Host
	if u.baseUrl.User != nil {
		if password, set := u.baseUrl.User.Password(); set {
			uploadReq.URL.User = u.baseUrl.User
			uploadReq.SetBasicAuth(u.baseUrl.User.Username(), password)
		}
	}
	uploadReq.URL.Path = primaryUrl.Path
	uploadReq.URL.RawQuery = primaryUrl.RawQuery

	rsp, err = u.do(uploadReq)
	return rsp, uploadReq.URL, err
}

func (u *httpUploader) do(req *http.Request) (*http.Response, error) {
	rsp, err := u.client.Do(req)
	if err != nil {
		return nil, err
	}

	switch rsp.StatusCode {
	case http.StatusOK, http.StatusCreated:
		return rsp, nil
	}

	respBody, _ := ioutil.ReadAll(rsp.Body)
	rsp.Body.Close()
	return rsp, fmt.Errorf("status code: %d\n%s", rsp.StatusCode, string(respBody))
}

// CLOUD CONTROLLER POLLING HELPERS
// The following code is not specific to uploading droplets.
// As soon as we talk to cloud controller in more than one place,
// this code should be extracted into a library

const (
	JOB_QUEUED   = "queued"
	JOB_RUNNING  = "running"
	JOB_FAILED   = "failed"
	JOB_FINISHED = "finished"
)

func (u *httpUploader) Poll(target *url.URL, res *http.Response, closeChan <-chan bool, interval time.Duration) error {
	body, err := u.parsePollingResponse(res)
	if err != nil {
		return err
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		switch body.Entity.Status {
		case JOB_QUEUED, JOB_RUNNING:
		case JOB_FINISHED:
			return nil
		case JOB_FAILED:
			return fmt.Errorf("upload job failed")
		default:
			return fmt.Errorf("unknown job status: %s", body.Entity.Status)
		}

		select {
		case <-ticker.C:
			pollingUrl, err := url.Parse(body.Metadata.Url)
			if err != nil {
				return err
			}

			if pollingUrl.Host == "" {
				pollingUrl.Scheme = target.Scheme
				pollingUrl.Host = target.Host
			}
			res, err := u.client.Get(pollingUrl.String())
			if err != nil {
				return err
			}

			body, err = u.parsePollingResponse(res)
			if err != nil {
				return err
			}
		case <-closeChan:
			return fmt.Errorf("upstream request was cancelled")
		}
	}
}

func (u *httpUploader) parsePollingResponse(res *http.Response) (pollingResponseBody, error) {
	body := pollingResponseBody{}
	err := json.NewDecoder(res.Body).Decode(&body)
	res.Body.Close()
	return body, err
}

type pollingResponseBody struct {
	Metadata struct {
		Url string
	}
	Entity struct {
		Status string
	}
}
<commit_msg>Forward md5 header in slow path upload as well [#82415708]<commit_after>package uploader

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"time"
)

type Uploader interface {
	Upload(u *url.URL, filename string, r *http.Request) (*http.Response, *url.URL, error)
	Poll(u *url.URL, res *http.Response, closeChan <-chan bool, interval time.Duration) error
}

type httpUploader struct {
	baseUrl *url.URL
	client  *http.Client
}

func New(baseUrl *url.URL, transport http.RoundTripper) Uploader {
	return &httpUploader{
		baseUrl: baseUrl,
		client: &http.Client{
			Transport: transport,
		},
	}
}

const contentMD5Header = "Content-MD5"

func (u *httpUploader) Upload(primaryUrl *url.URL, filename string, r *http.Request) (*http.Response, *url.URL, error) {
	if r.ContentLength <= 0 {
		return &http.Response{StatusCode: http.StatusLengthRequired}, nil, fmt.Errorf("Missing Content Length")
	}
	defer r.Body.Close()

	uploadReq, err := newMultipartRequestFromReader(r.ContentLength, r.Body, "upload[droplet]", filename)
	if err != nil {
		return nil, nil, err
	}

	uploadReq.Header.Set(contentMD5Header, r.Header.Get(contentMD5Header))

	// try the fast path
	uploadReq.URL = primaryUrl
	if primaryUrl.User != nil {
		if password, set := primaryUrl.User.Password(); set {
			uploadReq.SetBasicAuth(primaryUrl.User.Username(), password)
		}
	}

	rsp, err := u.do(uploadReq)
	if err == nil {
		return rsp, uploadReq.URL, err
	}

	// not a connect (dial) error
	var nestedErr error = err
	if urlErr, ok := err.(*url.Error); ok {
		nestedErr = urlErr.Err
	}
	if netErr, ok := nestedErr.(*net.OpError); !ok || netErr.Op != "dial" {
		return rsp, nil, err
	}

	// try the slow path
	uploadReq, err = newMultipartRequestFromReader(r.ContentLength, r.Body, "upload[droplet]", filename)
	if err != nil {
		return nil, nil, err
	}

	uploadReq.Header.Set(contentMD5Header, r.Header.Get(contentMD5Header))

	uploadReq.URL.Scheme = u.baseUrl.Scheme
	uploadReq.URL.Host = u.baseUrl.Host
	if u.baseUrl.User != nil {
		if password, set := u.baseUrl.User.Password(); set {
			uploadReq.URL.User = u.baseUrl.User
			uploadReq.SetBasicAuth(u.baseUrl.User.Username(), password)
		}
	}
	uploadReq.URL.Path = primaryUrl.Path
	uploadReq.URL.RawQuery = primaryUrl.RawQuery

	rsp, err = u.do(uploadReq)
	return rsp, uploadReq.URL, err
}

func (u *httpUploader) do(req *http.Request) (*http.Response, error) {
	rsp, err := u.client.Do(req)
	if err != nil {
		return nil, err
	}

	switch rsp.StatusCode {
	case http.StatusOK, http.StatusCreated:
		return rsp, nil
	}

	respBody, _ := ioutil.ReadAll(rsp.Body)
	rsp.Body.Close()
	return rsp, fmt.Errorf("status code: %d\n%s", rsp.StatusCode, string(respBody))
}

// CLOUD CONTROLLER POLLING HELPERS
// The following code is not specific to uploading droplets.
// As soon as we talk to cloud controller in more than one place,
// this code should be extracted into a library

const (
	JOB_QUEUED   = "queued"
	JOB_RUNNING  = "running"
	JOB_FAILED   = "failed"
	JOB_FINISHED = "finished"
)

func (u *httpUploader) Poll(target *url.URL, res *http.Response, closeChan <-chan bool, interval time.Duration) error {
	body, err := u.parsePollingResponse(res)
	if err != nil {
		return err
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		switch body.Entity.Status {
		case JOB_QUEUED, JOB_RUNNING:
		case JOB_FINISHED:
			return nil
		case JOB_FAILED:
			return fmt.Errorf("upload job failed")
		default:
			return fmt.Errorf("unknown job status: %s", body.Entity.Status)
		}

		select {
		case <-ticker.C:
			pollingUrl, err := url.Parse(body.Metadata.Url)
			if err != nil {
				return err
			}

			if pollingUrl.Host == "" {
				pollingUrl.Scheme = target.Scheme
				pollingUrl.Host = target.Host
			}
			res, err := u.client.Get(pollingUrl.String())
			if err != nil {
				return err
			}

			body, err = u.parsePollingResponse(res)
			if err != nil {
				return err
			}
		case <-closeChan:
			return fmt.Errorf("upstream request was cancelled")
		}
	}
}

func (u *httpUploader) parsePollingResponse(res *http.Response) (pollingResponseBody, error) {
	body := pollingResponseBody{}
	err := json.NewDecoder(res.Body).Decode(&body)
	res.Body.Close()
	return body, err
}

type pollingResponseBody struct {
	Metadata struct {
		Url string
	}
	Entity struct {
		Status string
	}
}
<|endoftext|>


<commit_before>package upstream

import (
	"bytes"
	"github.com/ngmoco/falcore"
	"io"
	"net"
	"net/http"
	"time"
)

type passThruReadCloser struct {
	io.Reader
	io.Closer
}

type Upstream struct {
	Transport *UpstreamTransport
	// Will ignore https on the incoming request and always upstream http
	ForceHttp bool
	// Ping URL Path-only for checking upness
	PingPath string
}

func NewUpstream(transport *UpstreamTransport) *Upstream {
	u := new(Upstream)
	u.Transport = transport
	return u
}

func (u *Upstream) FilterRequest(request *falcore.Request) (res *http.Response) {
	var err error
	req := request.HttpRequest

	// Force the upstream to use http
	if u.ForceHttp || req.URL.Scheme == "" {
		req.URL.Scheme = "http"
		req.URL.Host = req.Host
	}
	before := time.Now()
	req.Header.Set("Connection", "Keep-Alive")
	var upstrRes *http.Response
	upstrRes, err = u.Transport.transport.RoundTrip(req)
	diff := falcore.TimeDiff(before, time.Now())
	if err == nil {
		// Copy response over to new record.  Remove connection noise.  Add some sanity.
		res = falcore.SimpleResponse(req, upstrRes.StatusCode, nil, "")
		if upstrRes.ContentLength > 0 && upstrRes.Body != nil {
			res.ContentLength = upstrRes.ContentLength
			res.Body = upstrRes.Body
		} else if upstrRes.ContentLength == 0 && upstrRes.Body != nil {
			// Any bytes?
			var testBuf [1]byte
			n, _ := io.ReadFull(upstrRes.Body, testBuf[:])
			if n == 1 {
				// Yes there are.  Chunked it is.
				res.TransferEncoding = []string{"chunked"}
				res.ContentLength = -1
				rc := &passThruReadCloser{
					io.MultiReader(bytes.NewBuffer(testBuf[:]), upstrRes.Body),
					upstrRes.Body,
				}

				res.Body = rc
			}
		} else if upstrRes.Body != nil {
			res.Body = upstrRes.Body
			res.ContentLength = -1
			res.TransferEncoding = []string{"chunked"}
		}
		// Copy over headers with a few exceptions
		res.Header = make(http.Header)
		for hn, hv := range upstrRes.Header {
			switch hn {
			case "Content-Length":
			case "Connection":
			case "Transfer-Encoding":
			default:
				res.Header[hn] = hv
			}
		}
	} else {
		if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
			falcore.Error("%s Upstream Timeout error: %v", request.ID, err)
			res = falcore.SimpleResponse(req, 504, nil, "Gateway Timeout\n")
			request.CurrentStage.Status = 2 // Fail
		} else {
			falcore.Error("%s Upstream error: %v", request.ID, err)
			res = falcore.SimpleResponse(req, 502, nil, "Bad Gateway\n")
			request.CurrentStage.Status = 2 // Fail
		}
	}
	falcore.Debug("%s [%s] [%s] %s s=%d Time=%.4f", request.ID, req.Method, u.Transport.host, req.URL, res.StatusCode, diff)
	return
}

func (u *Upstream) ping() (up bool, ok bool) {
	if u.PingPath != "" {
		// the url must be syntactically valid for this to work but the host will be ignored because we
		// are overriding the connection always
		request, err := http.NewRequest("GET", "http://localhost"+u.PingPath, nil)
		request.Header.Set("Connection", "Keep-Alive") // not sure if this should be here for a ping
		if err != nil {
			falcore.Error("Bad Ping request: %v", err)
			return false, true
		}
		res, err := u.Transport.transport.RoundTrip(request)

		if err != nil {
			falcore.Error("Failed Ping to %v:%v: %v", u.Transport.host, u.Transport.port, err)
			return false, true
		} else {
			res.Body.Close()
		}
		if res.StatusCode == 200 {
			return true, true
		}
		falcore.Error("Failed Ping to %v:%v: %v", u.Transport.host, u.Transport.port, res.Status)
		// bad status
		return false, true
	}
	return false, false
}
<commit_msg>close upstream body if it has no length<commit_after>package upstream

import (
	"bytes"
	"github.com/ngmoco/falcore"
	"io"
	"net"
	"net/http"
	"time"
)

type passThruReadCloser struct {
	io.Reader
	io.Closer
}

type Upstream struct {
	Transport *UpstreamTransport
	// Will ignore https on the incoming request and always upstream http
	ForceHttp bool
	// Ping URL Path-only for checking upness
	PingPath string
}

func NewUpstream(transport *UpstreamTransport) *Upstream {
	u := new(Upstream)
	u.Transport = transport
	return u
}

func (u *Upstream) FilterRequest(request *falcore.Request) (res *http.Response) {
	var err error
	req := request.HttpRequest

	// Force the upstream to use http
	if u.ForceHttp || req.URL.Scheme == "" {
		req.URL.Scheme = "http"
		req.URL.Host = req.Host
	}
	before := time.Now()
	req.Header.Set("Connection", "Keep-Alive")
	var upstrRes *http.Response
	upstrRes, err = u.Transport.transport.RoundTrip(req)
	diff := falcore.TimeDiff(before, time.Now())
	if err == nil {
		// Copy response over to new record.  Remove connection noise.  Add some sanity.
		res = falcore.SimpleResponse(req, upstrRes.StatusCode, nil, "")
		if upstrRes.ContentLength > 0 && upstrRes.Body != nil {
			res.ContentLength = upstrRes.ContentLength
			res.Body = upstrRes.Body
		} else if upstrRes.ContentLength == 0 && upstrRes.Body != nil {
			// Any bytes?
			var testBuf [1]byte
			n, _ := io.ReadFull(upstrRes.Body, testBuf[:])
			if n == 1 {
				// Yes there are.  Chunked it is.
				res.TransferEncoding = []string{"chunked"}
				res.ContentLength = -1
				rc := &passThruReadCloser{
					io.MultiReader(bytes.NewBuffer(testBuf[:]), upstrRes.Body),
					upstrRes.Body,
				}

				res.Body = rc
			} else {
				// There was an error reading the body
				upstrRes.Body.Close()
				res.ContentLength = 0
				res.Body = nil
			}
		} else if upstrRes.Body != nil {
			res.Body = upstrRes.Body
			res.ContentLength = -1
			res.TransferEncoding = []string{"chunked"}
		}
		// Copy over headers with a few exceptions
		res.Header = make(http.Header)
		for hn, hv := range upstrRes.Header {
			switch hn {
			case "Content-Length":
			case "Connection":
			case "Transfer-Encoding":
			default:
				res.Header[hn] = hv
			}
		}
	} else {
		if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
			falcore.Error("%s Upstream Timeout error: %v", request.ID, err)
			res = falcore.SimpleResponse(req, 504, nil, "Gateway Timeout\n")
			request.CurrentStage.Status = 2 // Fail
		} else {
			falcore.Error("%s Upstream error: %v", request.ID, err)
			res = falcore.SimpleResponse(req, 502, nil, "Bad Gateway\n")
			request.CurrentStage.Status = 2 // Fail
		}
	}
	falcore.Debug("%s [%s] [%s] %s s=%d Time=%.4f", request.ID, req.Method, u.Transport.host, req.URL, res.StatusCode, diff)
	return
}

func (u *Upstream) ping() (up bool, ok bool) {
	if u.PingPath != "" {
		// the url must be syntactically valid for this to work but the host will be ignored because we
		// are overriding the connection always
		request, err := http.NewRequest("GET", "http://localhost"+u.PingPath, nil)
		request.Header.Set("Connection", "Keep-Alive") // not sure if this should be here for a ping
		if err != nil {
			falcore.Error("Bad Ping request: %v", err)
			return false, true
		}
		res, err := u.Transport.transport.RoundTrip(request)

		if err != nil {
			falcore.Error("Failed Ping to %v:%v: %v", u.Transport.host, u.Transport.port, err)
			return false, true
		} else {
			res.Body.Close()
		}
		if res.StatusCode == 200 {
			return true, true
		}
		falcore.Error("Failed Ping to %v:%v: %v", u.Transport.host, u.Transport.port, res.Status)
		// bad status
		return false, true
	}
	return false, false
}
<|endoftext|>