<commit_before>package core

import (
	"fmt"
	"strings"

	"github.com/af83/edwig/audit"
	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

type StopMonitoringRequestCollector interface {
	RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error)
}

type TestStopMonitoringRequestCollector struct {
	model.UUIDConsumer
}

type TestStopMonitoringRequestCollectorFactory struct{}

type SIRIStopMonitoringRequestCollector struct {
	model.ClockConsumer
	model.UUIDConsumer

	siriConnector
}

type SIRIStopMonitoringRequestCollectorFactory struct{}

func NewTestStopMonitoringRequestCollector() *TestStopMonitoringRequestCollector {
	return &TestStopMonitoringRequestCollector{}
}

// WIP
func (connector *TestStopMonitoringRequestCollector) RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error) {
	stopAreaUpdateEvent := model.NewStopAreaUpdateEvent(connector.NewUUID())
	stopAreaUpdateEvent.StopVisitUpdateEvents = append(stopAreaUpdateEvent.StopVisitUpdateEvents, &model.StopVisitUpdateEvent{})
	return stopAreaUpdateEvent, nil
}

func (factory *TestStopMonitoringRequestCollectorFactory) Validate(apiPartner *APIPartner) bool {
	return true
}

func (factory *TestStopMonitoringRequestCollectorFactory) CreateConnector(partner *Partner) Connector {
	return NewTestStopMonitoringRequestCollector()
}

func NewSIRIStopMonitoringRequestCollector(partner *Partner) *SIRIStopMonitoringRequestCollector {
	siriStopMonitoringRequestCollector := &SIRIStopMonitoringRequestCollector{}
	siriStopMonitoringRequestCollector.partner = partner
	return siriStopMonitoringRequestCollector
}

func (connector *SIRIStopMonitoringRequestCollector) RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error) {
	logStashEvent := make(audit.LogStashEvent)
	startTime := connector.Clock().Now()

	defer audit.CurrentLogStash().WriteEvent(logStashEvent)

	stopArea, ok := connector.Partner().Model().StopAreas().Find(request.StopAreaId())
	if !ok {
		return nil, fmt.Errorf("StopArea not found")
	}
	objectid, ok := stopArea.ObjectID(connector.partner.Setting("remote_objectid_kind"))
	if !ok {
		return nil, fmt.Errorf("StopArea %s doesn't have an ojbectID of type %s", stopArea.Id(), connector.partner.Setting("remote_objectid_kind"))
	}

	siriStopMonitoringRequest := &siri.SIRIStopMonitoringRequest{
		MessageIdentifier: connector.SIRIPartner().NewMessageIdentifier(),
		MonitoringRef:     objectid.Value(),
		RequestorRef:      connector.SIRIPartner().RequestorRef(),
		RequestTimestamp:  connector.Clock().Now(),
	}

	logSIRIStopMonitoringRequest(logStashEvent, siriStopMonitoringRequest)

	xmlStopMonitoringResponse, err := connector.SIRIPartner().SOAPClient().StopMonitoring(siriStopMonitoringRequest)
	logStashEvent["responseTime"] = connector.Clock().Since(startTime).String()
	if err != nil {
		logStashEvent["response"] = fmt.Sprintf("Error during CheckStatus: %v", err)
		return nil, err
	}

	logXMLStopMonitoringResponse(logStashEvent, xmlStopMonitoringResponse)

	// WIP
	stopAreaUpdateEvent := model.NewStopAreaUpdateEvent(connector.NewUUID())

	connector.setStopVisitUpdateEvents(stopAreaUpdateEvent, xmlStopMonitoringResponse)

	stopVisitsIds := []model.StopVisitId{}
	for _, stopVisit := range connector.Partner().Model().StopVisits().FindByStopAreaId(stopArea.Id()) {
		if stopVisit.IsCollected() == true {
			stopVisitsIds = append(stopVisitsIds, stopVisit.Id())
		}
	}

	connector.findAndSetStopVisitNotCollectedEvent(stopAreaUpdateEvent, stopVisitsIds)
	logStopVisitUpdateEvents(logStashEvent, stopAreaUpdateEvent)

	return stopAreaUpdateEvent, nil
}

func (connector *SIRIStopMonitoringRequestCollector) findAndSetStopVisitNotCollectedEvent(event *model.StopAreaUpdateEvent, stopVisitsIds []model.StopVisitId) {
	objId := make(map[string]bool)

	for _, stopVisitEvent := range event.StopVisitUpdateEvents {
		objId[stopVisitEvent.StopVisitObjectid.Value()] = true
	}

	for _, stopVisitid := range stopVisitsIds {
		if _, ok := objId[string(stopVisitid)]; !ok {
			stopVisitObjectId := model.NewObjectID(connector.partner.Setting("remote_objectid_kind"), string(stopVisitid))
			event.StopVisitNotCollectedEvents = append(event.StopVisitNotCollectedEvents, &model.StopVisitNotCollectedEvent{StopVisitObjectId: stopVisitObjectId})
		}
	}
}

func (connector *SIRIStopMonitoringRequestCollector) setStopVisitUpdateEvents(event *model.StopAreaUpdateEvent, xmlResponse *siri.XMLStopMonitoringResponse) {
	xmlStopVisitEvents := xmlResponse.XMLMonitoredStopVisits()
	if len(xmlStopVisitEvents) == 0 {
		return
	}

	for _, xmlStopVisitEvent := range xmlStopVisitEvents {
		stopVisitEvent := &model.StopVisitUpdateEvent{
			Id:                connector.NewUUID(),
			Created_at:        connector.Clock().Now(),
			RecordedAt:        xmlStopVisitEvent.RecordedAt(),
			VehicleAtStop:     xmlStopVisitEvent.VehicleAtStop(),
			StopVisitObjectid: model.NewObjectID(connector.partner.Setting("remote_objectid_kind"), xmlStopVisitEvent.ItemIdentifier()),
			StopAreaObjectId:  model.NewObjectID(connector.partner.Setting("remote_objectid_kind"), xmlStopVisitEvent.StopPointRef()),
			Schedules:         make(model.StopVisitSchedules),
			DepartureStatus:   model.StopVisitDepartureStatus(xmlStopVisitEvent.DepartureStatus()),
			ArrivalStatuts:    model.StopVisitArrivalStatus(xmlStopVisitEvent.ArrivalStatus()),
			Attributes:        NewSIRIStopVisitUpdateAttributes(xmlStopVisitEvent, connector.partner.Setting("remote_objectid_kind")),
		}
		stopVisitEvent.Schedules = model.NewStopVisitSchedules()
		if !xmlStopVisitEvent.AimedDepartureTime().IsZero() || !xmlStopVisitEvent.AimedArrivalTime().IsZero() {
			stopVisitEvent.Schedules.SetSchedule(model.STOP_VISIT_SCHEDULE_AIMED, xmlStopVisitEvent.AimedDepartureTime(), xmlStopVisitEvent.AimedArrivalTime())
		}
		if !xmlStopVisitEvent.ExpectedDepartureTime().IsZero() || !xmlStopVisitEvent.ExpectedArrivalTime().IsZero() {
			stopVisitEvent.Schedules.SetSchedule(model.STOP_VISIT_SCHEDULE_EXPECTED, xmlStopVisitEvent.ExpectedDepartureTime(), xmlStopVisitEvent.ExpectedArrivalTime())
		}
		if !xmlStopVisitEvent.ActualDepartureTime().IsZero() || !xmlStopVisitEvent.ActualArrivalTime().IsZero() {
			stopVisitEvent.Schedules.SetSchedule(model.STOP_VISIT_SCHEDULE_ACTUAL, xmlStopVisitEvent.ActualDepartureTime(), xmlStopVisitEvent.ActualArrivalTime())
		}
		event.StopVisitUpdateEvents = append(event.StopVisitUpdateEvents, stopVisitEvent)
	}
}

func (factory *SIRIStopMonitoringRequestCollectorFactory) Validate(apiPartner *APIPartner) bool {
	ok := apiPartner.ValidatePresenceOfSetting("remote_objectid_kind")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_url")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_credential")
	return ok
}

func (factory *SIRIStopMonitoringRequestCollectorFactory) CreateConnector(partner *Partner) Connector {
	return NewSIRIStopMonitoringRequestCollector(partner)
}

func logSIRIStopMonitoringRequest(logStashEvent audit.LogStashEvent, request *siri.SIRIStopMonitoringRequest) {
	logStashEvent["messageIdentifier"] = request.MessageIdentifier
	logStashEvent["monitoringRef"] = request.MonitoringRef
	logStashEvent["requestorRef"] = request.RequestorRef
	logStashEvent["requestTimestamp"] = request.RequestTimestamp.String()
	xml, err := request.BuildXML()
	if err != nil {
		logStashEvent["requestXML"] = fmt.Sprintf("%v", err)
		return
	}
	logStashEvent["requestXML"] = xml
}

func logXMLStopMonitoringResponse(logStashEvent audit.LogStashEvent, response *siri.XMLStopMonitoringResponse) {
	logStashEvent["address"] = response.Address()
	logStashEvent["producerRef"] = response.ProducerRef()
	logStashEvent["requestMessageRef"] = response.RequestMessageRef()
	logStashEvent["responseMessageIdentifier"] = response.ResponseMessageIdentifier()
	logStashEvent["responseTimestamp"] = response.ResponseTimestamp().String()
	logStashEvent["responseXML"] = response.RawXML()
}

func logStopVisitUpdateEvents(logStashEvent audit.LogStashEvent, stopAreaUpdateEvent *model.StopAreaUpdateEvent) {
	var idArray []string
	for _, stopVisitUpdateEvent := range stopAreaUpdateEvent.StopVisitUpdateEvents {
		idArray = append(idArray, stopVisitUpdateEvent.Id)
	}
	logStashEvent["StopVisitUpdateEventIds"] = strings.Join(idArray, ", ")
}
<commit_msg>Fixes SIRIStopMonitoringRequestCollector#findAndSetStopVisitNotCollectedEvent by using the StopVisit ObjectID associated to the connector. Refs #2828<commit_after>package core

import (
	"fmt"
	"strings"

	"github.com/af83/edwig/audit"
	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

type StopMonitoringRequestCollector interface {
	RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error)
}

type TestStopMonitoringRequestCollector struct {
	model.UUIDConsumer
}

type TestStopMonitoringRequestCollectorFactory struct{}

type SIRIStopMonitoringRequestCollector struct {
	model.ClockConsumer
	model.UUIDConsumer

	siriConnector
}

type SIRIStopMonitoringRequestCollectorFactory struct{}

func NewTestStopMonitoringRequestCollector() *TestStopMonitoringRequestCollector {
	return &TestStopMonitoringRequestCollector{}
}

// WIP
func (connector *TestStopMonitoringRequestCollector) RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error) {
	stopAreaUpdateEvent := model.NewStopAreaUpdateEvent(connector.NewUUID())
	stopAreaUpdateEvent.StopVisitUpdateEvents = append(stopAreaUpdateEvent.StopVisitUpdateEvents, &model.StopVisitUpdateEvent{})
	return stopAreaUpdateEvent, nil
}

func (factory *TestStopMonitoringRequestCollectorFactory) Validate(apiPartner *APIPartner) bool {
	return true
}

func (factory *TestStopMonitoringRequestCollectorFactory) CreateConnector(partner *Partner) Connector {
	return NewTestStopMonitoringRequestCollector()
}

func NewSIRIStopMonitoringRequestCollector(partner *Partner) *SIRIStopMonitoringRequestCollector {
	siriStopMonitoringRequestCollector := &SIRIStopMonitoringRequestCollector{}
	siriStopMonitoringRequestCollector.partner = partner
	return siriStopMonitoringRequestCollector
}

func (connector *SIRIStopMonitoringRequestCollector) RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error) {
	logStashEvent := make(audit.LogStashEvent)
	startTime := connector.Clock().Now()

	defer audit.CurrentLogStash().WriteEvent(logStashEvent)

	stopArea, ok := connector.Partner().Model().StopAreas().Find(request.StopAreaId())
	if !ok {
		return nil, fmt.Errorf("StopArea not found")
	}

	objectidKind := connector.partner.Setting("remote_objectid_kind")
	objectid, ok := stopArea.ObjectID(connector.partner.Setting("remote_objectid_kind"))
	if !ok {
		return nil, fmt.Errorf("StopArea %s doesn't have an ojbectID of type %s", stopArea.Id(), connector.partner.Setting("remote_objectid_kind"))
	}

	siriStopMonitoringRequest := &siri.SIRIStopMonitoringRequest{
		MessageIdentifier: connector.SIRIPartner().NewMessageIdentifier(),
		MonitoringRef:     objectid.Value(),
		RequestorRef:      connector.SIRIPartner().RequestorRef(),
		RequestTimestamp:  connector.Clock().Now(),
	}

	logSIRIStopMonitoringRequest(logStashEvent, siriStopMonitoringRequest)

	xmlStopMonitoringResponse, err := connector.SIRIPartner().SOAPClient().StopMonitoring(siriStopMonitoringRequest)
	logStashEvent["responseTime"] = connector.Clock().Since(startTime).String()
	if err != nil {
		logStashEvent["response"] = fmt.Sprintf("Error during CheckStatus: %v", err)
		return nil, err
	}

	logXMLStopMonitoringResponse(logStashEvent, xmlStopMonitoringResponse)

	// WIP
	stopAreaUpdateEvent := model.NewStopAreaUpdateEvent(connector.NewUUID())

	connector.setStopVisitUpdateEvents(stopAreaUpdateEvent, xmlStopMonitoringResponse)

	collectedStopVisitObjectIDs := []model.ObjectID{}
	for _, stopVisit := range connector.Partner().Model().StopVisits().FindByStopAreaId(stopArea.Id()) {
		if stopVisit.IsCollected() == true {
			objectId, ok := stopVisit.ObjectID(objectidKind)
			if ok {
				collectedStopVisitObjectIDs = append(collectedStopVisitObjectIDs, objectId)
			}
		}
	}

	connector.findAndSetStopVisitNotCollectedEvent(stopAreaUpdateEvent, collectedStopVisitObjectIDs)
	logStopVisitUpdateEvents(logStashEvent, stopAreaUpdateEvent)

	return stopAreaUpdateEvent, nil
}

func (connector *SIRIStopMonitoringRequestCollector) findAndSetStopVisitNotCollectedEvent(event *model.StopAreaUpdateEvent, collectedStopVisitObjectIDs []model.ObjectID) {
	objId := make(map[model.ObjectID]bool)

	for _, stopVisitEvent := range event.StopVisitUpdateEvents {
		objId[stopVisitEvent.StopVisitObjectid] = true
	}

	for _, stopVisitObjectID := range collectedStopVisitObjectIDs {
		if _, ok := objId[stopVisitObjectID]; !ok {
			event.StopVisitNotCollectedEvents = append(event.StopVisitNotCollectedEvents, &model.StopVisitNotCollectedEvent{StopVisitObjectId: stopVisitObjectID})
		}
	}
}

func (connector *SIRIStopMonitoringRequestCollector) setStopVisitUpdateEvents(event *model.StopAreaUpdateEvent, xmlResponse *siri.XMLStopMonitoringResponse) {
	xmlStopVisitEvents := xmlResponse.XMLMonitoredStopVisits()
	if len(xmlStopVisitEvents) == 0 {
		return
	}

	for _, xmlStopVisitEvent := range xmlStopVisitEvents {
		stopVisitEvent := &model.StopVisitUpdateEvent{
			Id:                connector.NewUUID(),
			Created_at:        connector.Clock().Now(),
			RecordedAt:        xmlStopVisitEvent.RecordedAt(),
			VehicleAtStop:     xmlStopVisitEvent.VehicleAtStop(),
			StopVisitObjectid: model.NewObjectID(connector.partner.Setting("remote_objectid_kind"), xmlStopVisitEvent.ItemIdentifier()),
			StopAreaObjectId:  model.NewObjectID(connector.partner.Setting("remote_objectid_kind"), xmlStopVisitEvent.StopPointRef()),
			Schedules:         make(model.StopVisitSchedules),
			DepartureStatus:   model.StopVisitDepartureStatus(xmlStopVisitEvent.DepartureStatus()),
			ArrivalStatuts:    model.StopVisitArrivalStatus(xmlStopVisitEvent.ArrivalStatus()),
			Attributes:        NewSIRIStopVisitUpdateAttributes(xmlStopVisitEvent, connector.partner.Setting("remote_objectid_kind")),
		}
		stopVisitEvent.Schedules = model.NewStopVisitSchedules()
		if !xmlStopVisitEvent.AimedDepartureTime().IsZero() || !xmlStopVisitEvent.AimedArrivalTime().IsZero() {
			stopVisitEvent.Schedules.SetSchedule(model.STOP_VISIT_SCHEDULE_AIMED, xmlStopVisitEvent.AimedDepartureTime(), xmlStopVisitEvent.AimedArrivalTime())
		}
		if !xmlStopVisitEvent.ExpectedDepartureTime().IsZero() || !xmlStopVisitEvent.ExpectedArrivalTime().IsZero() {
			stopVisitEvent.Schedules.SetSchedule(model.STOP_VISIT_SCHEDULE_EXPECTED, xmlStopVisitEvent.ExpectedDepartureTime(), xmlStopVisitEvent.ExpectedArrivalTime())
		}
		if !xmlStopVisitEvent.ActualDepartureTime().IsZero() || !xmlStopVisitEvent.ActualArrivalTime().IsZero() {
			stopVisitEvent.Schedules.SetSchedule(model.STOP_VISIT_SCHEDULE_ACTUAL, xmlStopVisitEvent.ActualDepartureTime(), xmlStopVisitEvent.ActualArrivalTime())
		}
		event.StopVisitUpdateEvents = append(event.StopVisitUpdateEvents, stopVisitEvent)
	}
}

func (factory *SIRIStopMonitoringRequestCollectorFactory) Validate(apiPartner *APIPartner) bool {
	ok := apiPartner.ValidatePresenceOfSetting("remote_objectid_kind")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_url")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_credential")
	return ok
}

func (factory *SIRIStopMonitoringRequestCollectorFactory) CreateConnector(partner *Partner) Connector {
	return NewSIRIStopMonitoringRequestCollector(partner)
}

func logSIRIStopMonitoringRequest(logStashEvent audit.LogStashEvent, request *siri.SIRIStopMonitoringRequest) {
	logStashEvent["messageIdentifier"] = request.MessageIdentifier
	logStashEvent["monitoringRef"] = request.MonitoringRef
	logStashEvent["requestorRef"] = request.RequestorRef
	logStashEvent["requestTimestamp"] = request.RequestTimestamp.String()
	xml, err := request.BuildXML()
	if err != nil {
		logStashEvent["requestXML"] = fmt.Sprintf("%v", err)
		return
	}
	logStashEvent["requestXML"] = xml
}

func logXMLStopMonitoringResponse(logStashEvent audit.LogStashEvent, response *siri.XMLStopMonitoringResponse) {
	logStashEvent["address"] = response.Address()
	logStashEvent["producerRef"] = response.ProducerRef()
	logStashEvent["requestMessageRef"] = response.RequestMessageRef()
	logStashEvent["responseMessageIdentifier"] = response.ResponseMessageIdentifier()
	logStashEvent["responseTimestamp"] = response.ResponseTimestamp().String()
	logStashEvent["responseXML"] = response.RawXML()
}

func logStopVisitUpdateEvents(logStashEvent audit.LogStashEvent, stopAreaUpdateEvent *model.StopAreaUpdateEvent) {
	var idArray []string
	for _, stopVisitUpdateEvent := range stopAreaUpdateEvent.StopVisitUpdateEvents {
		idArray = append(idArray, stopVisitUpdateEvent.Id)
	}
	logStashEvent["StopVisitUpdateEventIds"] = strings.Join(idArray, ", ")
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"

	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

type StopMonitoringRequestCollector interface {
	RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error)
}

type TestStopMonitoringRequestCollector struct {
	model.UUIDConsumer
}

type TestStopMonitoringRequestCollectorFactory struct{}

type SIRIStopMonitoringRequestCollector struct {
	model.ClockConsumer
	model.UUIDConsumer

	SIRIConnector

	objectid_kind string
}

type SIRIStopMonitoringRequestCollectorFactory struct{}

func NewTestStopMonitoringRequestCollector() *TestStopMonitoringRequestCollector {
	return &TestStopMonitoringRequestCollector{}
}

// WIP
func (connector *TestStopMonitoringRequestCollector) RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error) {
	stopAreaUpdateEvent := model.NewStopAreaUpdateEvent(connector.NewUUID())
	return stopAreaUpdateEvent, nil
}

func (factory *TestStopMonitoringRequestCollectorFactory) Validate(apiPartner *APIPartner) bool {
	return true
}

func (factory *TestStopMonitoringRequestCollectorFactory) CreateConnector(partner *Partner) Connector {
	return NewTestStopMonitoringRequestCollector()
}

func NewSIRIStopMonitoringRequestCollector(partner *Partner) *SIRIStopMonitoringRequestCollector {
	siriStopMonitoringRequestCollector := &SIRIStopMonitoringRequestCollector{
		objectid_kind: partner.Setting("remote_objectid_kind"),
	}
	siriStopMonitoringRequestCollector.partner = partner
	return siriStopMonitoringRequestCollector
}

func (connector *SIRIStopMonitoringRequestCollector) RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error) {
	stopArea, ok := connector.Partner().Model().StopAreas().Find(request.StopAreaId())
	if !ok {
		return nil, fmt.Errorf("StopArea not found")
	}
	objectid, ok := stopArea.ObjectID(connector.objectid_kind)
	if !ok {
		return nil, fmt.Errorf("stopArea doesn't have an ojbectID of type %s", connector.objectid_kind)
	}

	siriStopMonitoringRequest := &siri.SIRIStopMonitoringRequest{
		MessageIdentifier: connector.SIRIPartner().NewMessageIdentifier(),
		MonitoringRef:     objectid.Value(),
		RequestorRef:      connector.SIRIPartner().RequestorRef(),
		RequestTimestamp:  connector.Clock().Now(),
	}

	xmlStopMonitoringResponse, err := connector.SIRIPartner().SOAPClient().StopMonitoring(siriStopMonitoringRequest)
	if err != nil {
		return nil, err
	}

	// WIP
	stopAreaUpdateEvent := model.NewStopAreaUpdateEvent(connector.NewUUID())
	connector.setStopVisitUpdateEvents(stopAreaUpdateEvent, xmlStopMonitoringResponse)

	return stopAreaUpdateEvent, nil
}

func (connector *SIRIStopMonitoringRequestCollector) setStopVisitUpdateEvents(event *model.StopAreaUpdateEvent, xmlResponse *siri.XMLStopMonitoringResponse) {
	xmlStopVisitEvents := xmlResponse.XMLMonitoredStopVisits()
	if len(xmlStopVisitEvents) == 0 {
		return
	}
	for _, xmlStopVisitEvent := range xmlStopVisitEvents {
		stopVisitEvent := &model.StopVisitUpdateEvent{
			Id:                  connector.NewUUID(),
			Created_at:          connector.Clock().Now(),
			Stop_visit_objectid: model.NewObjectID(connector.objectid_kind, xmlStopVisitEvent.ItemIdentifier()),
			Schedules:           make(model.StopVisitSchedules),
			DepartureStatus:     model.StopVisitDepartureStatus(xmlStopVisitEvent.DepartureStatus()),
			ArrivalStatuts:      model.StopVisitArrivalStatus(xmlStopVisitEvent.ArrivalStatus()),
		}
		stopVisitEvent.Schedules = model.NewStopVisitSchedules()
		stopVisitEvent.SetSchedule(model.STOP_VISIT_SCHEDULE_AIMED, xmlStopVisitEvent.AimedDepartureTime(), xmlStopVisitEvent.AimedArrivalTime())
		stopVisitEvent.SetSchedule(model.STOP_VISIT_SCHEDULE_EXPECTED, xmlStopVisitEvent.ExpectedDepartureTime(), xmlStopVisitEvent.ExpectedArrivalTime())
		stopVisitEvent.SetSchedule(model.STOP_VISIT_SCHEDULE_ACTUAL, xmlStopVisitEvent.ActualDepartureTime(), xmlStopVisitEvent.ActualArrivalTime())
		event.StopVisitUpdateEvents = append(event.StopVisitUpdateEvents, stopVisitEvent)
	}
}

func (factory *SIRIStopMonitoringRequestCollectorFactory) Validate(apiPartner *APIPartner) bool {
	ok := apiPartner.ValidatePresenceOfSetting("remote_objectid_kind")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_url")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_credential")
	return ok
}

func (factory *SIRIStopMonitoringRequestCollectorFactory) CreateConnector(partner *Partner) Connector {
	return NewSIRIStopMonitoringRequestCollector(partner)
}
<commit_msg>Add StopArea id in SIRIStopMonitoringRequestCollector RequestStopAreaUpdate error<commit_after>package core

import (
	"fmt"

	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

type StopMonitoringRequestCollector interface {
	RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error)
}

type TestStopMonitoringRequestCollector struct {
	model.UUIDConsumer
}

type TestStopMonitoringRequestCollectorFactory struct{}

type SIRIStopMonitoringRequestCollector struct {
	model.ClockConsumer
	model.UUIDConsumer

	SIRIConnector

	objectid_kind string
}

type SIRIStopMonitoringRequestCollectorFactory struct{}

func NewTestStopMonitoringRequestCollector() *TestStopMonitoringRequestCollector {
	return &TestStopMonitoringRequestCollector{}
}

// WIP
func (connector *TestStopMonitoringRequestCollector) RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error) {
	stopAreaUpdateEvent := model.NewStopAreaUpdateEvent(connector.NewUUID())
	return stopAreaUpdateEvent, nil
}

func (factory *TestStopMonitoringRequestCollectorFactory) Validate(apiPartner *APIPartner) bool {
	return true
}

func (factory *TestStopMonitoringRequestCollectorFactory) CreateConnector(partner *Partner) Connector {
	return NewTestStopMonitoringRequestCollector()
}

func NewSIRIStopMonitoringRequestCollector(partner *Partner) *SIRIStopMonitoringRequestCollector {
	siriStopMonitoringRequestCollector := &SIRIStopMonitoringRequestCollector{
		objectid_kind: partner.Setting("remote_objectid_kind"),
	}
	siriStopMonitoringRequestCollector.partner = partner
	return siriStopMonitoringRequestCollector
}

func (connector *SIRIStopMonitoringRequestCollector) RequestStopAreaUpdate(request *StopAreaUpdateRequest) (*model.StopAreaUpdateEvent, error) {
	stopArea, ok := connector.Partner().Model().StopAreas().Find(request.StopAreaId())
	if !ok {
		return nil, fmt.Errorf("StopArea not found")
	}
	objectid, ok := stopArea.ObjectID(connector.objectid_kind)
	if !ok {
		return nil, fmt.Errorf("StopArea %s doesn't have an ojbectID of type %s", stopArea.Id(), connector.objectid_kind)
	}

	siriStopMonitoringRequest := &siri.SIRIStopMonitoringRequest{
		MessageIdentifier: connector.SIRIPartner().NewMessageIdentifier(),
		MonitoringRef:     objectid.Value(),
		RequestorRef:      connector.SIRIPartner().RequestorRef(),
		RequestTimestamp:  connector.Clock().Now(),
	}

	xmlStopMonitoringResponse, err := connector.SIRIPartner().SOAPClient().StopMonitoring(siriStopMonitoringRequest)
	if err != nil {
		return nil, err
	}

	// WIP
	stopAreaUpdateEvent := model.NewStopAreaUpdateEvent(connector.NewUUID())
	connector.setStopVisitUpdateEvents(stopAreaUpdateEvent, xmlStopMonitoringResponse)

	return stopAreaUpdateEvent, nil
}

func (connector *SIRIStopMonitoringRequestCollector) setStopVisitUpdateEvents(event *model.StopAreaUpdateEvent, xmlResponse *siri.XMLStopMonitoringResponse) {
	xmlStopVisitEvents := xmlResponse.XMLMonitoredStopVisits()
	if len(xmlStopVisitEvents) == 0 {
		return
	}
	for _, xmlStopVisitEvent := range xmlStopVisitEvents {
		stopVisitEvent := &model.StopVisitUpdateEvent{
			Id:                  connector.NewUUID(),
			Created_at:          connector.Clock().Now(),
			Stop_visit_objectid: model.NewObjectID(connector.objectid_kind, xmlStopVisitEvent.ItemIdentifier()),
			Schedules:           make(model.StopVisitSchedules),
			DepartureStatus:     model.StopVisitDepartureStatus(xmlStopVisitEvent.DepartureStatus()),
			ArrivalStatuts:      model.StopVisitArrivalStatus(xmlStopVisitEvent.ArrivalStatus()),
		}
		stopVisitEvent.Schedules = model.NewStopVisitSchedules()
		stopVisitEvent.SetSchedule(model.STOP_VISIT_SCHEDULE_AIMED, xmlStopVisitEvent.AimedDepartureTime(), xmlStopVisitEvent.AimedArrivalTime())
		stopVisitEvent.SetSchedule(model.STOP_VISIT_SCHEDULE_EXPECTED, xmlStopVisitEvent.ExpectedDepartureTime(), xmlStopVisitEvent.ExpectedArrivalTime())
		stopVisitEvent.SetSchedule(model.STOP_VISIT_SCHEDULE_ACTUAL, xmlStopVisitEvent.ActualDepartureTime(), xmlStopVisitEvent.ActualArrivalTime())
		event.StopVisitUpdateEvents = append(event.StopVisitUpdateEvents, stopVisitEvent)
	}
}

func (factory *SIRIStopMonitoringRequestCollectorFactory) Validate(apiPartner *APIPartner) bool {
	ok := apiPartner.ValidatePresenceOfSetting("remote_objectid_kind")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_url")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_credential")
	return ok
}

func (factory *SIRIStopMonitoringRequestCollectorFactory) CreateConnector(partner *Partner) Connector {
	return NewSIRIStopMonitoringRequestCollector(partner)
}
<|endoftext|>


<commit_before>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package downloader

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"

	"github.com/bogem/id3v2"
	"github.com/bogem/nehm/applescript"
	"github.com/bogem/nehm/config"
	"github.com/bogem/nehm/logs"
	"github.com/bogem/nehm/track"
	"github.com/valyala/fasthttp"
)

type Downloader struct {
	// dist is the folder, where tracks will be downloaded.
	dist string

	// itunesPlaylist is the iTunes playlist, where tracks will be added.
	itunesPlaylist string
}

func NewConfiguredDownloader() *Downloader {
	return &Downloader{
		dist:           config.Get("dlFolder"),
		itunesPlaylist: config.Get("itunesPlaylist"),
	}
}

func (downloader Downloader) DownloadAll(tracks []track.Track) {
	if len(tracks) == 0 {
		logs.FATAL.Println("there are no tracks to download")
	}

	var errors []string
	// Start with last track.
	for i := len(tracks) - 1; i >= 0; i-- {
		track := tracks[i]
		if err := downloader.Download(track); err != nil {
			logs.FEEDBACK.Println(" ✘")
			errors = append(errors, track.Fullname()+": "+err.Error())
			logs.ERROR.Printf("there was an error while downloading %v: %v\n\n", track.Fullname(), err)
		} else {
			logs.FEEDBACK.Println(" ✔")
		}
	}

	if len(errors) > 0 && len(tracks) > 1 {
		logs.FEEDBACK.Println("There were errors while downloading tracks:")
		for _, err := range errors {
			logs.FEEDBACK.Println("  " + err)
		}
		logs.FEEDBACK.Println()
	}
}

// artworks and trackBuf are used for reusing memory while downloading artworks and tracks.
var (
	artworkBuf []byte
	trackBuf   []byte
)

func (downloader Downloader) Download(t track.Track) error {
	// Create track file.
	trackPath := filepath.Join(downloader.dist, t.Filename())
	trackFile, e := os.Create(trackPath)
	if e != nil {
		return fmt.Errorf("couldn't create track file: %v", e)
	}

	// err lets us to not prevent the processing of track further.
	// err will only be returned at the end of this function.
	var err error

	logs.INFO.Printf("Downloading from %q to %q\n", t.URL(), trackPath)
	logs.INFO.Printf("Downloading artwork from %q\n", t.ArtworkURL())
	logs.FEEDBACK.Printf("Downloading %q ...", t.Fullname())

	// Parallelize downloading of track and artwork.
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		// Download artwork.
		artworkBuf = artworkBuf[:0]
		_, artworkBuf, e = fasthttp.Get(artworkBuf, t.ArtworkURL())
		if e != nil {
			err = fmt.Errorf("couldn't download artwork file: %v", e)
		}
		// Write ID3 tag to trackFile.
		if e := writeTagToWriter(t, trackFile, artworkBuf); e != nil && err == nil {
			err = fmt.Errorf("there was an error while tagging track: %v", e)
		}
	}()

	// Download track.
	trackBuf = trackBuf[:0]
	_, trackBuf, e = fasthttp.Get(trackBuf, t.URL())
	if e != nil {
		return fmt.Errorf("couldn't download track: %v", e)
	}

	wg.Wait()

	// Write track to track file.
	if _, e := trackFile.Write(trackBuf); e != nil {
		return fmt.Errorf("couldn't write track to file: %v", e)
	}

	// Add to iTunes.
	if downloader.itunesPlaylist != "" {
		logs.FEEDBACK.Print(" adding to iTunes ...")
		if e := applescript.AddTrackToPlaylist(trackPath, downloader.itunesPlaylist); e != nil && err == nil {
			err = fmt.Errorf("couldn't add track to playlist: %v", e)
		}
	}

	return err
}

func writeTagToWriter(t track.Track, w io.Writer, artwork []byte) error {
	tag := id3v2.NewEmptyTag()

	tag.SetArtist(t.Artist())
	tag.SetTitle(t.Title())
	tag.SetYear(t.Year())

	if len(artwork) > 0 {
		pic := id3v2.PictureFrame{
			Encoding:    id3v2.ENUTF8,
			MimeType:    "image/jpeg",
			PictureType: id3v2.PTFrontCover,
			Picture:     artwork,
		}
		tag.AddAttachedPicture(pic)
	}

	_, err := tag.WriteTo(w)
	return err
}
<commit_msg>downloader: Prevent writing tag to track file after the error in downloading<commit_after>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package downloader

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"

	"github.com/bogem/id3v2"
	"github.com/bogem/nehm/applescript"
	"github.com/bogem/nehm/config"
	"github.com/bogem/nehm/logs"
	"github.com/bogem/nehm/track"
	"github.com/valyala/fasthttp"
)

type Downloader struct {
	// dist is the folder, where tracks will be downloaded.
	dist string

	// itunesPlaylist is the iTunes playlist, where tracks will be added.
	itunesPlaylist string
}

func NewConfiguredDownloader() *Downloader {
	return &Downloader{
		dist:           config.Get("dlFolder"),
		itunesPlaylist: config.Get("itunesPlaylist"),
	}
}

func (downloader Downloader) DownloadAll(tracks []track.Track) {
	if len(tracks) == 0 {
		logs.FATAL.Println("there are no tracks to download")
	}

	var errors []string
	// Start with last track.
	for i := len(tracks) - 1; i >= 0; i-- {
		track := tracks[i]
		if err := downloader.Download(track); err != nil {
			logs.FEEDBACK.Println(" ✘")
			errors = append(errors, track.Fullname()+": "+err.Error())
			logs.ERROR.Printf("there was an error while downloading %v: %v\n\n", track.Fullname(), err)
		} else {
			logs.FEEDBACK.Println(" ✔")
		}
	}

	if len(errors) > 0 && len(tracks) > 1 {
		logs.FEEDBACK.Println("There were errors while downloading tracks:")
		for _, err := range errors {
			logs.FEEDBACK.Println("  " + err)
		}
		logs.FEEDBACK.Println()
	}
}

// artworks and trackBuf are used for reusing memory while downloading artworks and tracks.
var (
	artworkBuf []byte
	trackBuf   []byte
)

func (downloader Downloader) Download(t track.Track) error {
	// Create track file.
	trackPath := filepath.Join(downloader.dist, t.Filename())
	trackFile, e := os.Create(trackPath)
	if e != nil {
		return fmt.Errorf("couldn't create track file: %v", e)
	}

	// err lets us to not prevent the processing of track further.
	// err will only be returned at the end of this function.
	var err error

	logs.INFO.Printf("Downloading from %q to %q\n", t.URL(), trackPath)
	logs.INFO.Printf("Downloading artwork from %q\n", t.ArtworkURL())
	logs.FEEDBACK.Printf("Downloading %q ...", t.Fullname())

	// Parallelize downloading of track and artwork.
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()

		// Download artwork.
		artworkBuf = artworkBuf[:0]
		_, artworkBuf, e = fasthttp.Get(artworkBuf, t.ArtworkURL())
		if e != nil {
			err = fmt.Errorf("couldn't download artwork file: %v", e)
			return
		}

		// Write ID3 tag to trackFile.
		if e := writeTagToWriter(t, trackFile, artworkBuf); e != nil {
			err = fmt.Errorf("there was an error while tagging track: %v", e)
		}
	}()

	// Download track.
	trackBuf = trackBuf[:0]
	_, trackBuf, e = fasthttp.Get(trackBuf, t.URL())
	if e != nil {
		return fmt.Errorf("couldn't download track: %v", e)
	}

	wg.Wait()

	// Write track to track file.
	if _, e := trackFile.Write(trackBuf); e != nil {
		return fmt.Errorf("couldn't write track to file: %v", e)
	}

	// Add to iTunes.
	if downloader.itunesPlaylist != "" {
		logs.FEEDBACK.Print(" adding to iTunes ...")
		if e := applescript.AddTrackToPlaylist(trackPath, downloader.itunesPlaylist); e != nil && err == nil {
			err = fmt.Errorf("couldn't add track to playlist: %v", e)
		}
	}

	return err
}

func writeTagToWriter(t track.Track, w io.Writer, artwork []byte) error {
	tag := id3v2.NewEmptyTag()

	tag.SetArtist(t.Artist())
	tag.SetTitle(t.Title())
	tag.SetYear(t.Year())

	if len(artwork) > 0 {
		pic := id3v2.PictureFrame{
			Encoding:    id3v2.ENUTF8,
			MimeType:    "image/jpeg",
			PictureType: id3v2.PTFrontCover,
			Picture:     artwork,
		}
		tag.AddAttachedPicture(pic)
	}

	_, err := tag.WriteTo(w)
	return err
}
<|endoftext|>


<commit_before>package rabbit

import (
	"log"
	"time"

	"github.com/streadway/amqp"
)

const unconfirmedMessagesMaxCount = 1000

//AssuredPublisher allows you to publish events to RabbitMQ with implicit delivery confirmation
type AssuredPublisher struct {
	Publisher

	unconfirmedMessages     map[uint64]*unconfirmedMessage
	sequenceNumber          uint64
	waitAfterEachPublishing bool
	closeChannel            chan *amqp.Error
	confirmationHandler     func(amqp.Confirmation, interface{})
	doNotRepublish          bool
}

type unconfirmedMessage struct {
	message    []byte
	subscriber *Subscriber
	arg        interface{}
}

func (p *AssuredPublisher) construct() {
	p.initNewChannel()
	p.NotifyPublish(unconfirmedMessagesMaxCount)
	for err := p.Confirm(false); err != nil; err = p.Confirm(false) {
		logError(err, "Can't setup confirmations for a publisher")
		select {
		case <-p.cancelChannel:
			return
		case <-time.After(1 * time.Second):
		}
	}
}

// should be within locked block
func (p *AssuredPublisher) initNewChannel() {
	channel := p.getChannelWithoutLock()
	closeChannel := channel.NotifyClose(make(chan *amqp.Error, 1))
	p.lock.Lock()
	p.closeChannel = closeChannel
	p.sequenceNumber = 0
	p.lock.Unlock()
	go p.continuoslyReceiveConfirmations()
}

// should be within locked block
func (p *AssuredPublisher) ensureChannel() {
	p.getChannelWithoutLock()
	p.lock.Lock()
	closeChannel := p.closeChannel
	p.lock.Unlock()
	select {
	case e := <-closeChannel:
		log.Printf("Error: publisher's channel has been closed (%+#v). Recreating it", e)
		p.lock.Lock()
		p._channel = nil
		p.lock.Unlock()
		p.reconnect()
	default:
	}
}

// should be within locked block
func (p *AssuredPublisher) reconnect() bool {
	p.receiveAllConfirmations()
	p.initNewChannel()
	return p.republishAllMessages()
}

func (p *AssuredPublisher) republishAllMessages() bool {
	messages := map[uint64]*unconfirmedMessage{}
	p.lock.Lock()
	for deliveryTag, message := range p.unconfirmedMessages {
		messages[deliveryTag] = message
	}
	p.unconfirmedMessages = map[uint64]*unconfirmedMessage{0: &unconfirmedMessage{}}
	doNotRepublish := p.doNotRepublish
	confirmationHandler := p.confirmationHandler
	p.lock.Unlock()
	for deliveryTag, message := range messages {
		if deliveryTag == 0 {
			continue
		}
		if !doNotRepublish {
			if !p.publishBytesWithArgWithoutLock(message.message, message.subscriber, message.arg) { // if cancelled
				return false
			}
		} else {
			if confirmationHandler != nil {
				confirmationHandler(amqp.Confirmation{Ack: false, DeliveryTag: deliveryTag}, message.arg)
			}
		}
	}
	p.lock.Lock()
	delete(p.unconfirmedMessages, 0)
	p.lock.Unlock()
	return true
}

// NewAssuredPublisher constructs a new AssuredPublisher instance
func NewAssuredPublisher(cancelChannel <-chan bool) *AssuredPublisher {
	publisher := &AssuredPublisher{Publisher: Publisher{connection: publisherConnection, cancelChannel: cancelChannel},
		unconfirmedMessages: map[uint64]*unconfirmedMessage{}, waitAfterEachPublishing: true}
	publisher.construct()
	return publisher
}

// NewAssuredPublisherWithConnection constructs a new AssuredPublisher instance
func NewAssuredPublisherWithConnection(connection *Connection, cancelChannel <-chan bool) *AssuredPublisher {
	publisher := &AssuredPublisher{Publisher: Publisher{connection: connection, cancelChannel: cancelChannel},
		unconfirmedMessages: map[uint64]*unconfirmedMessage{}, waitAfterEachPublishing: true}
	publisher.construct()
	return publisher
}

// SetExplicitWaiting disables implicit waiting for a confirmation after each publishing
func (p *AssuredPublisher) SetExplicitWaiting() {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	p.lock.Lock()
	p.waitAfterEachPublishing = false
	p.lock.Unlock()
}

// DisableRepublishing disables messages republishing
func (p *AssuredPublisher) DisableRepublishing() {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	p.lock.Lock()
	p.doNotRepublish = true
	p.lock.Unlock()
}

// SetConfirmationHandler sets the handler which is called for every confirmation received
func (p *AssuredPublisher) SetConfirmationHandler(confirmationHandler func(amqp.Confirmation, interface{})) {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	p.lock.Lock()
	p.confirmationHandler = confirmationHandler
	p.lock.Unlock()
}

// Publish pushes items on to a RabbitMQ Queue.
// For AssuredPublisher it waits for delivery confirmaiton and retries on failures
func (p *AssuredPublisher) Publish(message string, subscriber *Subscriber) bool {
	return p.PublishBytes([]byte(message), subscriber)
}

// PublishWithArg pushes items on to a RabbitMQ Queue. The argument will be stored for passing into the confirmation handler.
// For AssuredPublisher it waits for delivery confirmaiton and retries on failures
func (p *AssuredPublisher) PublishWithArg(message string, subscriber *Subscriber, arg interface{}) bool {
	return p.PublishBytesWithArg([]byte(message), subscriber, arg)
}

// PublishBytes is the same as Publish but accepts a []byte instead of a string.
// For AssuredPublisher it waits for delivery confirmaiton and retries on failures
func (p *AssuredPublisher) PublishBytes(message []byte, subscriber *Subscriber) bool {
	return p.PublishBytesWithArg(message, subscriber, nil)
}

// PublishBytesWithArg is the same as Publish but accepts a []byte instead of a string.
// The argument will be stored for passing into the confirmation handler.
// For AssuredPublisher it waits for delivery confirmaiton and retries on failures
func (p *AssuredPublisher) PublishBytesWithArg(message []byte, subscriber *Subscriber, arg interface{}) bool {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	return p.publishBytesWithArgWithoutLock(message, subscriber, arg)
}

func (p *AssuredPublisher) publishBytesWithArgWithoutLock(message []byte, subscriber *Subscriber, arg interface{}) bool {
	for {
		p.ensureChannel()
		p.receiveAllConfirmations()
		p.lock.Lock()
		unconfirmedCount := len(p.unconfirmedMessages)
		p.lock.Unlock()
		if unconfirmedCount >= unconfirmedMessagesMaxCount {
			if !p.waitForConfirmation() {
				return false
			}
		}
		if err := (&p.Publisher).publishBytesWithoutLock(message, subscriber); err != nil {
			log.Printf("Error on pushing into RabbitMQ: %v", err)
			select {
			case <-p.cancelChannel:
				return false
			case <-time.After(1 * time.Second):
			}
			continue
		}
		break
	}
	p.lock.Lock()
	p.sequenceNumber++
	var body []byte
	if !p.doNotRepublish {
		body = message
	}
	p.unconfirmedMessages[p.sequenceNumber] = &unconfirmedMessage{body, subscriber, arg}
	p.lock.Unlock()
	if p.waitAfterEachPublishing && !p.waitForConfirmation() {
		return false
	}
	return true
}

func (p *AssuredPublisher) waitForConfirmation() bool {
	for {
		p.lock.Lock()
		ch := p._notifyPublish[0].channel
		p.lock.Unlock()
		select {
		case confirm := <-ch:
			if !p.handleConfirmationMessage(confirm) {
				p.lock.Lock()
				if len(p.unconfirmedMessages) == 0 {
					p.lock.Unlock()
					return true
				}
				p.lock.Unlock()
			} else {
				return true
			}
		case <-time.After(10 * time.Second):
			log.Printf("Error: RabbitMQ Timeout")
		case <-p.cancelChannel:
			return false
		}
		log.Printf("Next in waitForConfirmation")
	}
}

// ReceiveAllConfirmations explicitly receives all awaiting confirmations
func (p *AssuredPublisher) ReceiveAllConfirmations() bool {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	return p.receiveAllConfirmations()
}

func (p *AssuredPublisher) continuoslyReceiveConfirmations() {
	for {
		if !p.receiveAllConfirmations() {
			return
		}
		select {
		case <-p.cancelChannel:
			return
		case <-time.After(1 * time.Second):
		}
	}
}

func (p *AssuredPublisher) receiveAllConfirmations() bool {
	p.lock.Lock()
	count := len(p.unconfirmedMessages)
	if len(p._notifyPublish) < 1 {
		p.lock.Unlock()
		return true
	}
	channel := p._notifyPublish[0].channel
	p.lock.Unlock()
	if count == 0 {
		return true
	}
	for {
		select {
		case confirm := <-channel:
			if !p.handleConfirmationMessage(confirm) {
				return false
			}
		default:
			return true
		}
	}
}

func (p *AssuredPublisher) handleConfirmationMessage(confirm amqp.Confirmation) bool {
	if confirm.DeliveryTag == 0 {
		log.Printf("Unknown Error (RabbitMQ Publisher's confirmation channel has been closed)")
		return false
	}
	p.lock.Lock()
	confirmationHandler := p.confirmationHandler
	arg := p.unconfirmedMessages[confirm.DeliveryTag].arg
	p.lock.Unlock()
	if confirmationHandler != nil {
		confirmationHandler(confirm, arg)
	}
	if confirm.Ack || p.doNotRepublish {
		p.forgetMessage(confirm.DeliveryTag)
	}
	if !confirm.Ack {
		log.Printf("Unknown Error (RabbitMQ Ack is false)")
	}
	return true
}

// forgetMessage removes a message from the internal storage
func (p *AssuredPublisher) forgetMessage(deliveryTag uint64) {
	p.lock.Lock()
	delete(p.unconfirmedMessages, deliveryTag)
	p.lock.Unlock()
}

// WaitForAllConfirmations waits for all confirmations and retries publishing if needed.
// Returns false only if is cancelled.
func (p *AssuredPublisher) WaitForAllConfirmations() bool {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	for {
		p.lock.Lock()
		if len(p.unconfirmedMessages) == 0 {
			p.lock.Unlock()
			return true
		}
		p.lock.Unlock()
		select {
		case <-p.cancelChannel:
			return false
		case <-time.After(300 * time.Millisecond):
		}
	}
}
<commit_msg>tune delays (#28)<commit_after>package rabbit

import (
	"log"
	"time"

	"github.com/streadway/amqp"
)

const unconfirmedMessagesMaxCount = 1000

//AssuredPublisher allows you to publish events to RabbitMQ with implicit delivery confirmation
type AssuredPublisher struct {
	Publisher

	unconfirmedMessages     map[uint64]*unconfirmedMessage
	sequenceNumber          uint64
	waitAfterEachPublishing bool
	closeChannel            chan *amqp.Error
	confirmationHandler     func(amqp.Confirmation, interface{})
	doNotRepublish          bool
}

type unconfirmedMessage struct {
	message    []byte
	subscriber *Subscriber
	arg        interface{}
}

func (p *AssuredPublisher) construct() {
	p.initNewChannel()
	p.NotifyPublish(unconfirmedMessagesMaxCount)
	for err := p.Confirm(false); err != nil; err = p.Confirm(false) {
		logError(err, "Can't setup confirmations for a publisher")
		select {
		case <-p.cancelChannel:
			return
		case <-time.After(1 * time.Second):
		}
	}
}

// should be within locked block
func (p *AssuredPublisher) initNewChannel() {
	channel := p.getChannelWithoutLock()
	closeChannel := channel.NotifyClose(make(chan *amqp.Error, 1))
	p.lock.Lock()
	p.closeChannel = closeChannel
	p.sequenceNumber = 0
	p.lock.Unlock()
	go p.continuoslyReceiveConfirmations()
}

// should be within locked block
func (p *AssuredPublisher) ensureChannel() {
	p.getChannelWithoutLock()
	p.lock.Lock()
	closeChannel := p.closeChannel
	p.lock.Unlock()
	select {
	case e := <-closeChannel:
		log.Printf("Error: publisher's channel has been closed (%+#v). Recreating it", e)
		p.lock.Lock()
		p._channel = nil
		p.lock.Unlock()
		p.reconnect()
	default:
	}
}

// should be within locked block
func (p *AssuredPublisher) reconnect() bool {
	p.receiveAllConfirmations()
	p.initNewChannel()
	return p.republishAllMessages()
}

func (p *AssuredPublisher) republishAllMessages() bool {
	messages := map[uint64]*unconfirmedMessage{}
	p.lock.Lock()
	for deliveryTag, message := range p.unconfirmedMessages {
		messages[deliveryTag] = message
	}
	p.unconfirmedMessages = map[uint64]*unconfirmedMessage{0: &unconfirmedMessage{}}
	doNotRepublish := p.doNotRepublish
	confirmationHandler := p.confirmationHandler
	p.lock.Unlock()
	for deliveryTag, message := range messages {
		if deliveryTag == 0 {
			continue
		}
		if !doNotRepublish {
			if !p.publishBytesWithArgWithoutLock(message.message, message.subscriber, message.arg) { // if cancelled
				return false
			}
		} else {
			if confirmationHandler != nil {
				confirmationHandler(amqp.Confirmation{Ack: false, DeliveryTag: deliveryTag}, message.arg)
			}
		}
	}
	p.lock.Lock()
	delete(p.unconfirmedMessages, 0)
	p.lock.Unlock()
	return true
}

// NewAssuredPublisher constructs a new AssuredPublisher instance
func NewAssuredPublisher(cancelChannel <-chan bool) *AssuredPublisher {
	publisher := &AssuredPublisher{Publisher: Publisher{connection: publisherConnection, cancelChannel: cancelChannel},
		unconfirmedMessages: map[uint64]*unconfirmedMessage{}, waitAfterEachPublishing: true}
	publisher.construct()
	return publisher
}

// NewAssuredPublisherWithConnection constructs a new AssuredPublisher instance
func NewAssuredPublisherWithConnection(connection *Connection, cancelChannel <-chan bool) *AssuredPublisher {
	publisher := &AssuredPublisher{Publisher: Publisher{connection: connection, cancelChannel: cancelChannel},
		unconfirmedMessages: map[uint64]*unconfirmedMessage{}, waitAfterEachPublishing: true}
	publisher.construct()
	return publisher
}

// SetExplicitWaiting disables implicit waiting for a confirmation after each publishing
func (p *AssuredPublisher) SetExplicitWaiting() {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	p.lock.Lock()
	p.waitAfterEachPublishing = false
	p.lock.Unlock()
}

// DisableRepublishing disables messages republishing
func (p *AssuredPublisher) DisableRepublishing() {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	p.lock.Lock()
	p.doNotRepublish = true
	p.lock.Unlock()
}

// SetConfirmationHandler sets the handler which is called for every confirmation received
func (p *AssuredPublisher) SetConfirmationHandler(confirmationHandler func(amqp.Confirmation, interface{})) {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	p.lock.Lock()
	p.confirmationHandler = confirmationHandler
	p.lock.Unlock()
}

// Publish pushes items on to a RabbitMQ Queue.
// For AssuredPublisher it waits for delivery confirmaiton and retries on failures
func (p *AssuredPublisher) Publish(message string, subscriber *Subscriber) bool {
	return p.PublishBytes([]byte(message), subscriber)
}

// PublishWithArg pushes items on to a RabbitMQ Queue. The argument will be stored for passing into the confirmation handler.
// For AssuredPublisher it waits for delivery confirmaiton and retries on failures
func (p *AssuredPublisher) PublishWithArg(message string, subscriber *Subscriber, arg interface{}) bool {
	return p.PublishBytesWithArg([]byte(message), subscriber, arg)
}

// PublishBytes is the same as Publish but accepts a []byte instead of a string.
// For AssuredPublisher it waits for delivery confirmaiton and retries on failures
func (p *AssuredPublisher) PublishBytes(message []byte, subscriber *Subscriber) bool {
	return p.PublishBytesWithArg(message, subscriber, nil)
}

// PublishBytesWithArg is the same as Publish but accepts a []byte instead of a string.
// The argument will be stored for passing into the confirmation handler.
// For AssuredPublisher it waits for delivery confirmaiton and retries on failures
func (p *AssuredPublisher) PublishBytesWithArg(message []byte, subscriber *Subscriber, arg interface{}) bool {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	return p.publishBytesWithArgWithoutLock(message, subscriber, arg)
}

func (p *AssuredPublisher) publishBytesWithArgWithoutLock(message []byte, subscriber *Subscriber, arg interface{}) bool {
	for {
		p.ensureChannel()
		p.receiveAllConfirmations()
		p.lock.Lock()
		unconfirmedCount := len(p.unconfirmedMessages)
		p.lock.Unlock()
		if unconfirmedCount >= unconfirmedMessagesMaxCount {
			if !p.waitForConfirmation() {
				return false
			}
		}
		if err := (&p.Publisher).publishBytesWithoutLock(message, subscriber); err != nil {
			log.Printf("Error on pushing into RabbitMQ: %v", err)
			select {
			case <-p.cancelChannel:
				return false
			case <-time.After(1 * time.Second):
			}
			continue
		}
		break
	}
	p.lock.Lock()
	p.sequenceNumber++
	var body []byte
	if !p.doNotRepublish {
		body = message
	}
	p.unconfirmedMessages[p.sequenceNumber] = &unconfirmedMessage{body, subscriber, arg}
	p.lock.Unlock()
	if p.waitAfterEachPublishing && !p.waitForConfirmation() {
		return false
	}
	return true
}

func (p *AssuredPublisher) waitForConfirmation() bool {
	for {
		p.lock.Lock()
		ch := p._notifyPublish[0].channel
		p.lock.Unlock()
		select {
		case confirm := <-ch:
			if !p.handleConfirmationMessage(confirm) {
				p.lock.Lock()
				if len(p.unconfirmedMessages) == 0 {
					p.lock.Unlock()
					return true
				}
				p.lock.Unlock()
			} else {
				return true
			}
		case <-time.After(10 * time.Second):
			log.Printf("Error: RabbitMQ Timeout")
		case <-p.cancelChannel:
			return false
		}
		log.Printf("Next in waitForConfirmation")
	}
}

// ReceiveAllConfirmations explicitly receives all awaiting confirmations
func (p *AssuredPublisher) ReceiveAllConfirmations() bool {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	return p.receiveAllConfirmations()
}

func (p *AssuredPublisher) continuoslyReceiveConfirmations() {
	for {
		if !p.receiveAllConfirmations() {
			return
		}
		select {
		case <-p.cancelChannel:
			return
		case <-time.After(50 * time.Millisecond):
		}
	}
}

func (p *AssuredPublisher) receiveAllConfirmations() bool {
	p.lock.Lock()
	count := len(p.unconfirmedMessages)
	if len(p._notifyPublish) < 1 {
		p.lock.Unlock()
		return true
	}
	channel := p._notifyPublish[0].channel
	p.lock.Unlock()
	if count == 0 {
		return true
	}
	for {
		select {
		case confirm := <-channel:
			if !p.handleConfirmationMessage(confirm) {
				return false
			}
		default:
			return true
		}
	}
}

func (p *AssuredPublisher) handleConfirmationMessage(confirm amqp.Confirmation) bool {
	if confirm.DeliveryTag == 0 {
		log.Printf("Unknown Error (RabbitMQ Publisher's confirmation channel has been closed)")
		return false
	}
	p.lock.Lock()
	confirmationHandler := p.confirmationHandler
	arg := p.unconfirmedMessages[confirm.DeliveryTag].arg
	p.lock.Unlock()
	if confirmationHandler != nil {
		confirmationHandler(confirm, arg)
	}
	if confirm.Ack || p.doNotRepublish {
		p.forgetMessage(confirm.DeliveryTag)
	}
	if !confirm.Ack {
		log.Printf("Unknown Error (RabbitMQ Ack is false)")
	}
	return true
}

// forgetMessage removes a message from the internal storage
func (p *AssuredPublisher) forgetMessage(deliveryTag uint64) {
	p.lock.Lock()
	delete(p.unconfirmedMessages, deliveryTag)
	p.lock.Unlock()
}

// WaitForAllConfirmations waits for all confirmations and retries publishing if needed.
// Returns false only if is cancelled.
func (p *AssuredPublisher) WaitForAllConfirmations() bool {
	p.publicMethodsLock.Lock()
	defer p.publicMethodsLock.Unlock()
	for {
		p.lock.Lock()
		if len(p.unconfirmedMessages) == 0 {
			p.lock.Unlock()
			return true
		}
		p.lock.Unlock()
		select {
		case <-p.cancelChannel:
			return false
		case <-time.After(50 * time.Millisecond):
		}
	}
}
<|endoftext|>


<commit_before>// Package v3action contains the business logic for the commands/v3 package
package v3action

// This is used for sorting.
type SortOrder string

const (
	Ascending  SortOrder = "Ascending"
	Descending SortOrder = "Descending"
)

// Warnings is a list of warnings returned back from the cloud controller
type Warnings []string

// Actor represents a V3 actor.
type Actor struct {
	CloudControllerClient CloudControllerClient
	Config                Config
	SharedActor           SharedActor
	UAAClient             UAAClient
}

// NewActor returns a new V3 actor.
func NewActor(client CloudControllerClient, config Config, sharedActor SharedActor, uaaClient UAAClient) *Actor {
	return &Actor{
		CloudControllerClient: client,
		Config:                config,
		SharedActor:           sharedActor,
		UAAClient:             uaaClient,
	}
}
<commit_msg>fix comment<commit_after>// Package v3action contains the business logic for the commands/v3 package
package v3action

// SortOrder is used for sorting.
type SortOrder string

const (
	Ascending  SortOrder = "Ascending"
	Descending SortOrder = "Descending"
)

// Warnings is a list of warnings returned back from the cloud controller
type Warnings []string

// Actor represents a V3 actor.
type Actor struct {
	CloudControllerClient CloudControllerClient
	Config                Config
	SharedActor           SharedActor
	UAAClient             UAAClient
}

// NewActor returns a new V3 actor.
func NewActor(client CloudControllerClient, config Config, sharedActor SharedActor, uaaClient UAAClient) *Actor {
	return &Actor{
		CloudControllerClient: client,
		Config:                config,
		SharedActor:           sharedActor,
		UAAClient:             uaaClient,
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package base64

import (
	"base64";
	"bytes";
	"io";
	"os";
	"reflect";
	"strings";
	"testing";
)

type testpair struct {
	decoded, encoded string;
}

var pairs = []testpair {
	// RFC 3548 examples
	testpair{"\x14\xfb\x9c\x03\xd9\x7e", "FPucA9l+"},
	testpair{"\x14\xfb\x9c\x03\xd9", "FPucA9k="},
	testpair{"\x14\xfb\x9c\x03", "FPucAw=="},

	// RFC 4648 examples
	testpair{"", ""},
	testpair{"f", "Zg=="},
	testpair{"fo", "Zm8="},
	testpair{"foo", "Zm9v"},
	testpair{"foob", "Zm9vYg=="},
	testpair{"fooba", "Zm9vYmE="},
	testpair{"foobar", "Zm9vYmFy"},

	// Wikipedia examples
	testpair{"sure.", "c3VyZS4="},
	testpair{"sure", "c3VyZQ=="},
	testpair{"sur", "c3Vy"},
	testpair{"su", "c3U="},
	testpair{"leasure.", "bGVhc3VyZS4="},
	testpair{"easure.", "ZWFzdXJlLg=="},
	testpair{"asure.", "YXN1cmUu"},
	testpair{"sure.", "c3VyZS4="},
}

var bigtest = testpair {
	"Twas brillig, and the slithy toves",
	"VHdhcyBicmlsbGlnLCBhbmQgdGhlIHNsaXRoeSB0b3Zlcw=="
}

func testEqual(t *testing.T, msg string, args ...) bool {
	v := reflect.NewValue(args).(reflect.StructValue);
	v1 := v.Field(v.Len() - 2);
	v2 := v.Field(v.Len() - 1);
	if v1.Interface() != v2.Interface() {
		t.Errorf(msg, args);
		return false;
	}
	return true;
}

func TestEncode(t *testing.T) {
	for _, p := range pairs {
		buf := make([]byte, StdEncoding.EncodedLen(len(p.decoded)));
		StdEncoding.Encode(strings.Bytes(p.decoded), buf);
		testEqual(t, "Encode(%q) = %q, want %q", p.decoded, string(buf), p.encoded);
	}
}

func TestEncoder(t *testing.T) {
	for _, p := range pairs {
		bb := &bytes.Buffer{};
		encoder := NewEncoder(StdEncoding, bb);
		encoder.Write(strings.Bytes(p.decoded));
		encoder.Close();
		testEqual(t, "Encode(%q) = %q, want %q", p.decoded, string(bb.Data()), p.encoded);
	}
}

func TestEncoderBuffering(t *testing.T) {
	input := strings.Bytes(bigtest.decoded);
	for bs := 1; bs <= 12; bs++ {
		buf := make([]byte, bs);
		bb := &bytes.Buffer{};
		encoder := NewEncoder(StdEncoding, bb);
		for pos := 0; pos < len(input); pos += bs {
			end := pos+bs;
			if end > len(input) {
				end = len(input);
			}
			n, err := encoder.Write(input[pos:end]);
			testEqual(t, "Write(%q) gave error %v, want %v", input[pos:end], err, os.Error(nil));
			testEqual(t, "Write(%q) gave length %v, want %v", input[pos:end], n, end-pos);
		}
		err := encoder.Close();
		testEqual(t, "Close gave error %v, want %v", err, os.Error(nil));
		testEqual(t, "Encoding/%d of %q = %q, want %q", bs, bigtest.decoded, string(bb.Data()), bigtest.encoded);
	}
}

func TestDecode(t *testing.T) {
	for _, p := range pairs {
		dbuf := make([]byte, StdEncoding.DecodedLen(len(p.encoded)));
		count, end, err := StdEncoding.decode(strings.Bytes(p.encoded), dbuf);
		testEqual(t, "Decode(%q) = error %v, want %v", p.encoded, err, os.Error(nil));
		testEqual(t, "Decode(%q) = length %v, want %v", p.encoded, count, len(p.decoded));
		if len(p.encoded) > 0 {
			testEqual(t, "Decode(%q) = end %v, want %v", p.encoded, end, (p.encoded[len(p.encoded)-1] == '='));
		}
		testEqual(t, "Decode(%q) = %q, want %q", p.encoded, string(dbuf[0:count]), p.decoded);
	}
}

func TestDecoder(t *testing.T) {
	for _, p := range pairs {
		decoder := NewDecoder(StdEncoding, io.NewByteReader(strings.Bytes(p.encoded)));
		dbuf := make([]byte, StdEncoding.DecodedLen(len(p.encoded)));
		count, err := decoder.Read(dbuf);
		if err != nil && err != os.EOF {
			t.Fatal("Read failed", err);
		}
		testEqual(t, "Read from %q = length %v, want %v", p.encoded, count, len(p.decoded));
		testEqual(t, "Decoding of %q = %q, want %q", p.encoded, string(dbuf[0:count]), p.decoded);
		if err != os.EOF {
			count, err = decoder.Read(dbuf);
		}
		testEqual(t, "Read from %q = %v, want %v", p.encoded, err, os.EOF);
	}
}

func TestDecoderBuffering(t *testing.T) {
	input := strings.Bytes(bigtest.encoded);
	for bs := 1; bs <= 12; bs++ {
		decoder := NewDecoder(StdEncoding, io.NewByteReader(input));
		buf := make([]byte, len(bigtest.decoded) + 12);
		var total int;
		for total = 0; total < len(bigtest.decoded); {
			n, err := decoder.Read(buf[total:total+bs]);
			testEqual(t, "Read from %q at pos %d = %d, %v, want _, %v", bigtest.encoded, total, n, err, os.Error(nil));
			total += n;
		}
		testEqual(t, "Decoding/%d of %q = %q, want %q", bs, bigtest.encoded, string(buf[0:total]), bigtest.decoded);
	}
}

func TestDecodeCorrupt(t *testing.T) {
	type corrupt struct {
		e string;
		p int;
	};
	examples := []corrupt {
		corrupt{"!!!!", 0},
		corrupt{"x===", 1},
		corrupt{"AA=A", 2},
		corrupt{"AAA=AAAA", 3},
		corrupt{"AAAAA", 4},
		corrupt{"AAAAAA", 4},
	};

	for _, e := range examples {
		dbuf := make([]byte, StdEncoding.DecodedLen(len(e.e)));
		count, err := StdEncoding.Decode(strings.Bytes(e.e), dbuf);
		switch err := err.(type) {
		case CorruptInputError:
			testEqual(t, "Corruption in %q at offset %v, want %v", e.e, int(err), e.p);
		default:
			t.Error("Decoder failed to detect corruption in", e);
		}
	}
}

func TestBig(t *testing.T) {
	n := 3*1000+1;
	raw := make([]byte, n);
	const alpha = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
	for i := 0; i < n; i++ {
		raw[i] = alpha[i%len(alpha)];
	}
	encoded := new(bytes.Buffer);
	w := NewEncoder(StdEncoding, encoded);
	nn, err := w.Write(raw);
	if nn != n || err != nil {
		t.Fatalf("Encoder.Write(raw) = %d, %v want %d, nil", nn, err, n);
	}
	err = w.Close();
	if err != nil {
		t.Fatalf("Encoder.Close() = %v want nil", err);
	}
	decoded, err := io.ReadAll(NewDecoder(StdEncoding, encoded));
	if err != nil {
		t.Fatalf("io.ReadAll(NewDecoder(...)): %v", err);
	}

	if !bytes.Equal(raw, decoded) {
		var i int;
		for i = 0; i < len(decoded) && i < len(raw); i++ {
			if decoded[i] != raw[i] {
				break;
			}
		}
		t.Errorf("Decode(Encode(%d-byte string)) failed at offset %d", n, i);
	}
}
<commit_msg>base64_test: use new reflect interface (CL 31107)<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package base64

import (
	"base64";
	"bytes";
	"io";
	"os";
	"reflect";
	"strings";
	"testing";
)

type testpair struct {
	decoded, encoded string;
}

var pairs = []testpair {
	// RFC 3548 examples
	testpair{"\x14\xfb\x9c\x03\xd9\x7e", "FPucA9l+"},
	testpair{"\x14\xfb\x9c\x03\xd9", "FPucA9k="},
	testpair{"\x14\xfb\x9c\x03", "FPucAw=="},

	// RFC 4648 examples
	testpair{"", ""},
	testpair{"f", "Zg=="},
	testpair{"fo", "Zm8="},
	testpair{"foo", "Zm9v"},
	testpair{"foob", "Zm9vYg=="},
	testpair{"fooba", "Zm9vYmE="},
	testpair{"foobar", "Zm9vYmFy"},

	// Wikipedia examples
	testpair{"sure.", "c3VyZS4="},
	testpair{"sure", "c3VyZQ=="},
	testpair{"sur", "c3Vy"},
	testpair{"su", "c3U="},
	testpair{"leasure.", "bGVhc3VyZS4="},
	testpair{"easure.", "ZWFzdXJlLg=="},
	testpair{"asure.", "YXN1cmUu"},
	testpair{"sure.", "c3VyZS4="},
}

var bigtest = testpair {
	"Twas brillig, and the slithy toves",
	"VHdhcyBicmlsbGlnLCBhbmQgdGhlIHNsaXRoeSB0b3Zlcw=="
}

func testEqual(t *testing.T, msg string, args ...) bool {
	v := reflect.NewValue(args).(*reflect.StructValue);
	v1 := v.Field(v.NumField() - 2);
	v2 := v.Field(v.NumField() - 1);
	if v1.Interface() != v2.Interface() {
		t.Errorf(msg, args);
		return false;
	}
	return true;
}

func TestEncode(t *testing.T) {
	for _, p := range pairs {
		buf := make([]byte, StdEncoding.EncodedLen(len(p.decoded)));
		StdEncoding.Encode(strings.Bytes(p.decoded), buf);
		testEqual(t, "Encode(%q) = %q, want %q", p.decoded, string(buf), p.encoded);
	}
}

func TestEncoder(t *testing.T) {
	for _, p := range pairs {
		bb := &bytes.Buffer{};
		encoder := NewEncoder(StdEncoding, bb);
		encoder.Write(strings.Bytes(p.decoded));
		encoder.Close();
		testEqual(t, "Encode(%q) = %q, want %q", p.decoded, string(bb.Data()), p.encoded);
	}
}

func TestEncoderBuffering(t *testing.T) {
	input := strings.Bytes(bigtest.decoded);
	for bs := 1; bs <= 12; bs++ {
		buf := make([]byte, bs);
		bb := &bytes.Buffer{};
		encoder := NewEncoder(StdEncoding, bb);
		for pos := 0; pos < len(input); pos += bs {
			end := pos+bs;
			if end > len(input) {
				end = len(input);
			}
			n, err := encoder.Write(input[pos:end]);
			testEqual(t, "Write(%q) gave error %v, want %v", input[pos:end], err, os.Error(nil));
			testEqual(t, "Write(%q) gave length %v, want %v", input[pos:end], n, end-pos);
		}
		err := encoder.Close();
		testEqual(t, "Close gave error %v, want %v", err, os.Error(nil));
		testEqual(t, "Encoding/%d of %q = %q, want %q", bs, bigtest.decoded, string(bb.Data()), bigtest.encoded);
	}
}

func TestDecode(t *testing.T) {
	for _, p := range pairs {
		dbuf := make([]byte, StdEncoding.DecodedLen(len(p.encoded)));
		count, end, err := StdEncoding.decode(strings.Bytes(p.encoded), dbuf);
		testEqual(t, "Decode(%q) = error %v, want %v", p.encoded, err, os.Error(nil));
		testEqual(t, "Decode(%q) = length %v, want %v", p.encoded, count, len(p.decoded));
		if len(p.encoded) > 0 {
			testEqual(t, "Decode(%q) = end %v, want %v", p.encoded, end, (p.encoded[len(p.encoded)-1] == '='));
		}
		testEqual(t, "Decode(%q) = %q, want %q", p.encoded, string(dbuf[0:count]), p.decoded);
	}
}

func TestDecoder(t *testing.T) {
	for _, p := range pairs {
		decoder := NewDecoder(StdEncoding, io.NewByteReader(strings.Bytes(p.encoded)));
		dbuf := make([]byte, StdEncoding.DecodedLen(len(p.encoded)));
		count, err := decoder.Read(dbuf);
		if err != nil && err != os.EOF {
			t.Fatal("Read failed", err);
		}
		testEqual(t, "Read from %q = length %v, want %v", p.encoded, count, len(p.decoded));
		testEqual(t, "Decoding of %q = %q, want %q", p.encoded, string(dbuf[0:count]), p.decoded);
		if err != os.EOF {
			count, err = decoder.Read(dbuf);
		}
		testEqual(t, "Read from %q = %v, want %v", p.encoded, err, os.EOF);
	}
}

func TestDecoderBuffering(t *testing.T) {
	input := strings.Bytes(bigtest.encoded);
	for bs := 1; bs <= 12; bs++ {
		decoder := NewDecoder(StdEncoding, io.NewByteReader(input));
		buf := make([]byte, len(bigtest.decoded) + 12);
		var total int;
		for total = 0; total < len(bigtest.decoded); {
			n, err := decoder.Read(buf[total:total+bs]);
			testEqual(t, "Read from %q at pos %d = %d, %v, want _, %v", bigtest.encoded, total, n, err, os.Error(nil));
			total += n;
		}
		testEqual(t, "Decoding/%d of %q = %q, want %q", bs, bigtest.encoded, string(buf[0:total]), bigtest.decoded);
	}
}

func TestDecodeCorrupt(t *testing.T) {
	type corrupt struct {
		e string;
		p int;
	};
	examples := []corrupt {
		corrupt{"!!!!", 0},
		corrupt{"x===", 1},
		corrupt{"AA=A", 2},
		corrupt{"AAA=AAAA", 3},
		corrupt{"AAAAA", 4},
		corrupt{"AAAAAA", 4},
	};

	for _, e := range examples {
		dbuf := make([]byte, StdEncoding.DecodedLen(len(e.e)));
		count, err := StdEncoding.Decode(strings.Bytes(e.e), dbuf);
		switch err := err.(type) {
		case CorruptInputError:
			testEqual(t, "Corruption in %q at offset %v, want %v", e.e, int(err), e.p);
		default:
			t.Error("Decoder failed to detect corruption in", e);
		}
	}
}

func TestBig(t *testing.T) {
	n := 3*1000+1;
	raw := make([]byte, n);
	const alpha = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
	for i := 0; i < n; i++ {
		raw[i] = alpha[i%len(alpha)];
	}
	encoded := new(bytes.Buffer);
	w := NewEncoder(StdEncoding, encoded);
	nn, err := w.Write(raw);
	if nn != n || err != nil {
		t.Fatalf("Encoder.Write(raw) = %d, %v want %d, nil", nn, err, n);
	}
	err = w.Close();
	if err != nil {
		t.Fatalf("Encoder.Close() = %v want nil", err);
	}
	decoded, err := io.ReadAll(NewDecoder(StdEncoding, encoded));
	if err != nil {
		t.Fatalf("io.ReadAll(NewDecoder(...)): %v", err);
	}

	if !bytes.Equal(raw, decoded) {
		var i int;
		for i = 0; i < len(decoded) && i < len(raw); i++ {
			if decoded[i] != raw[i] {
				break;
			}
		}
		t.Errorf("Decode(Encode(%d-byte string)) failed at offset %d", n, i);
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Pagoda Box Inc
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not distributed with this file, You can obtain one
// at http://mozilla.org/MPL/2.0/.
//

// pgconfig.go provides methods for configuring a node corresponding to if it is
// running a 'master' or 'slave' instance of postgres.

package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
)

//
type pgConfig struct {
	listenAddr string
	master     bool
}

// configureHBAConf attempts to open the 'pg_hba.conf' file. Once open it will scan
// the file line by line looking for replication settings, and overwrite only those
// settings with the settings required for redundancy on Pagoda Box
func configureHBAConf() error {

	// get the role of the other node in the cluster that is running an instance
	// of postgresql
	self := Whoami()
	other, err := Whoisnot(self.CRole)
	if err != nil {
		log.Warn("[pg_config.configureHBAConf] Unable to find another!\n%s\n", err)
	}

	// open the pg_hba.conf
	file := conf.DataDir + "pg_hba.conf"
	f, err := os.Open(file)
	if err != nil {
		log.Error("[pg_config.configureHBAConf] Failed to open '%s'!\n%s\n", file, err)
		return err
	}

	defer f.Close()

	scanner := bufio.NewScanner(f)
	entry := ""

	// scan the file line by line to build an 'entry' to be re-written back to the
	// file, skipping ('removing') any line that deals with redundancy.
	for scanner.Scan() {

		// stop scanning if a special prefix is encountered.
		if strings.HasPrefix(scanner.Text(), "#~") {
			break
		}

		// dont care about submatches, just if the string matches, 'skipping' any lines
		// that are custom configurations
		if regexp.MustCompile(`^\s*#?\s*(local|host)\s*(replication)`).FindString(scanner.Text()) == "" {
			entry += fmt.Sprintf("%s\n", scanner.Text())
		}
	}

	// if the other node is present, write redundancy into the 'entry' (otherwise
	// just leave it out)

	if other != nil {
		entry += fmt.Sprintf(`#~-----------------------------------------------------------------------------
# YOKE CONFIG
#------------------------------------------------------------------------------

# these configuration options have been removed from their standard location and
# placed here so that Pagoda Box could override them with the neccessary values
# to configure redundancy.

# IMPORTANT: these settings will always be overriden when the server boots. They
# are set dynamically and so should never change.

host    replication     %s        %s/32            trust`, SystemUser(), other.Ip)
	}

	// write the 'entry' to the file
	err = ioutil.WriteFile(file, []byte(entry), 0644)
	if err != nil {
		log.Error("[pg_config.configureHBAConf] Failed to write to '%s'!\n%s\n", file, err)
		return err
	}

	return nil
}

// configurePGConf attempts to open the 'postgresql.conf' file. Once open it will
// scan the file line by line looking for replication settings, and overwrite only
// those settings with the settings required for redundancy on Pagoda Box
func configurePGConf(opts pgConfig) error {

	// open the postgresql.conf
	file := conf.DataDir + "postgresql.conf"
	f, err := os.Open(file)
	if err != nil {
		return err
	}

	defer f.Close()

	scanner := bufio.NewScanner(f)
	entry := ""

	// scan the file line by line to build an 'entry' to be re-written back to the
	// file, skipping ('removing') any lines that need to be manually configured
	for scanner.Scan() {

		submatch := regexp.MustCompile(`^\s*#?\s*(listen_addresses|port|wal_level|archive_mode|archive_command|max_wal_senders|wal_keep_segments|hot_standby|synchronous_standby_names)\s*=\s*`).FindStringSubmatch(scanner.Text())

		// stop scanning if a special prefix is encountered. This ensures there are
		// no duplicate Pagoda Box comment blocks
		if strings.HasPrefix(scanner.Text(), "#~") {
			break
		}

		// build the 'entry' from all lines that don't match the custom configurations
		if submatch == nil {
			entry += fmt.Sprintf("%s\n", scanner.Text())
		}
	}

	// write manual configurations into an 'entry'
	entry += fmt.Sprintf(`#~-----------------------------------------------------------------------------
# YOKE CONFIG
#------------------------------------------------------------------------------

# these configuration options have been removed from their standard location and
# placed here so that Pagoda Box could override them with the neccessary values
# to configure redundancy.

# IMPORTANT: these settings will always be overriden when the server boots. They
# are set dynamically and so should never change.

listen_addresses = '%s'      # what IP address(es) to listen on;
                                  # comma-separated list of addresses;
                                  # defaults to 'localhost'; use '*' for all
                                  # (change requires restart)
port = %d                     # (change requires restart)
wal_level = hot_standby           # minimal, archive, or hot_standby
                                  # (change requires restart)
archive_mode = on                 # allows archiving to be done
                                  # (change requires restart)
archive_command = 'exit 0'        # command to use to archive a logfile segment
                                  # placeholders: \%p = path of file to archive
                                  #               \%f = file name only
                                  # e.g. 'test ! -f /mnt/server/archivedir/\%f && cp \%p /mnt/server/archivedir/\%f'
max_wal_senders = 10              # max number of walsender processes
                                  # (change requires restart)
wal_keep_segments = 16          	# in logfile segments, 16MB each; 0 disables
hot_standby = on                  # "on" allows queries during recovery
                                  # (change requires restart)
`, opts.listenAddr, conf.PGPort)

	// if this node is currenty 'master' then write one additional configuration
	// into the 'entry'
	if opts.master {
		entry += `
# added only the the node running postgres as 'master'
synchronous_standby_names = slave # standby servers that provide sync rep
                                  # comma-separated list of application_name
                                  # from standby(s); '*' = all
`
	}

	// write 'entry' to the file
	err = ioutil.WriteFile(file, []byte(entry), 0644)
	if err != nil {
		log.Error("[pg_config.configurePGConf] Failed to write to '%s'!\n%s\n", file, err)
	}

	return nil
}

// createRecovery creates a 'recovery.conf' file with the necessary settings
// required for redundancy on Pagoda Box. This method is called on the node that
// is being configured to run the 'slave' instance of postgres
func createRecovery() error {

	file := conf.DataDir + "recovery.conf"

	self := Whoami()
	other, err := Whoisnot(self.CRole)
	if err != nil {
		log.Fatal("[pg_config.createRecovery] Unable to find another... Exiting!\n%s\n", err)
		os.Exit(1)
	}

	// open/truncate the recover.conf
	f, err := os.Create(file)
	if err != nil {
		log.Error("[pg_config.createRecovery] Failed to create '%s'!\n%s\n", file, err)
		return err
	}

	// write manual configuration an 'entry'
	entry := fmt.Sprintf(`#~-----------------------------------------------------------------------------
# YOKE CONFIG
#------------------------------------------------------------------------------

# IMPORTANT: this config file is dynamically generated by Pagoda Box for redundancy
# any changes made here will be overriden.

# When standby_mode is enabled, the PostgreSQL server will work as a standby. It
# tries to connect to the primary according to the connection settings
# primary_conninfo, and receives XLOG records continuously.
standby_mode = on
primary_conninfo = 'host=%s port=%d application_name=slave'

# restore_command specifies the shell command that is executed to copy log files
# back from archival storage. This parameter is *required* for an archive
# recovery, but optional for streaming replication. The given command satisfies
# the requirement without doing anything.
restore_command = 'exit 0'
`, other.Ip, other.PGPort)

	// write 'entry' to the file
	if _, err := f.WriteString(entry); err != nil {
		log.Error("[pg_config.createRecovery] Failed to write to '%s'!\n%s\n", file, err)
		return err
	}

	return nil
}

// destroyRecovery attempts to destroy the 'recovery.conf'. This method is called
// on a node that is being configured to run the 'master' instance of postgres
func destroyRecovery() {

	file := conf.DataDir + "recovery.conf"

	// remove 'recovery.conf'
	err := os.Remove(file)
	if err != nil {
		log.Warn("[pg_config.destroyRecovery] No recovery.conf found at '%s'", file)
	}
}
<commit_msg>quick output update<commit_after>// Copyright (c) 2015 Pagoda Box Inc
//
// This Source Code Form is subject to the terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not distributed with this file, You can obtain one
// at http://mozilla.org/MPL/2.0/.
//

// pgconfig.go provides methods for configuring a node corresponding to if it is
// running a 'master' or 'slave' instance of postgres.

package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
)

//
type pgConfig struct {
	listenAddr string
	master     bool
}

// configureHBAConf attempts to open the 'pg_hba.conf' file. Once open it will scan
// the file line by line looking for replication settings, and overwrite only those
// settings with the settings required for redundancy on Pagoda Box
func configureHBAConf() error {

	// get the role of the other node in the cluster that is running an instance
	// of postgresql
	self := Whoami()
	other, err := Whoisnot(self.CRole)
	if err != nil {
		log.Warn("[pgconfig.configureHBAConf] Unable to find another!\n%s\n", err)
	}

	// open the pg_hba.conf
	file := conf.DataDir + "pg_hba.conf"
	f, err := os.Open(file)
	if err != nil {
		log.Error("[pgconfig.configureHBAConf] Failed to open '%s'!\n%s\n", file, err)
		return err
	}

	defer f.Close()

	scanner := bufio.NewScanner(f)
	entry := ""

	// scan the file line by line to build an 'entry' to be re-written back to the
	// file, skipping ('removing') any line that deals with redundancy.
	for scanner.Scan() {

		// stop scanning if a special prefix is encountered.
		if strings.HasPrefix(scanner.Text(), "#~") {
			break
		}

		// dont care about submatches, just if the string matches, 'skipping' any lines
		// that are custom configurations
		if regexp.MustCompile(`^\s*#?\s*(local|host)\s*(replication)`).FindString(scanner.Text()) == "" {
			entry += fmt.Sprintf("%s\n", scanner.Text())
		}
	}

	// if the other node is present, write redundancy into the 'entry' (otherwise
	// just leave it out)

	if other != nil {
		entry += fmt.Sprintf(`#~-----------------------------------------------------------------------------
# YOKE CONFIG
#------------------------------------------------------------------------------

# these configuration options have been removed from their standard location and
# placed here so that Pagoda Box could override them with the neccessary values
# to configure redundancy.

# IMPORTANT: these settings will always be overriden when the server boots. They
# are set dynamically and so should never change.

host    replication     %s        %s/32            trust`, SystemUser(), other.Ip)
	}

	// write the 'entry' to the file
	err = ioutil.WriteFile(file, []byte(entry), 0644)
	if err != nil {
		log.Error("[pgconfig.configureHBAConf] Failed to write to '%s'!\n%s\n", file, err)
		return err
	}

	return nil
}

// configurePGConf attempts to open the 'postgresql.conf' file. Once open it will
// scan the file line by line looking for replication settings, and overwrite only
// those settings with the settings required for redundancy on Pagoda Box
func configurePGConf(opts pgConfig) error {

	// open the postgresql.conf
	file := conf.DataDir + "postgresql.conf"
	f, err := os.Open(file)
	if err != nil {
		return err
	}

	defer f.Close()

	scanner := bufio.NewScanner(f)
	entry := ""

	// scan the file line by line to build an 'entry' to be re-written back to the
	// file, skipping ('removing') any lines that need to be manually configured
	for scanner.Scan() {

		submatch := regexp.MustCompile(`^\s*#?\s*(listen_addresses|port|wal_level|archive_mode|archive_command|max_wal_senders|wal_keep_segments|hot_standby|synchronous_standby_names)\s*=\s*`).FindStringSubmatch(scanner.Text())

		// stop scanning if a special prefix is encountered. This ensures there are
		// no duplicate Pagoda Box comment blocks
		if strings.HasPrefix(scanner.Text(), "#~") {
			break
		}

		// build the 'entry' from all lines that don't match the custom configurations
		if submatch == nil {
			entry += fmt.Sprintf("%s\n", scanner.Text())
		}
	}

	// write manual configurations into an 'entry'
	entry += fmt.Sprintf(`#~-----------------------------------------------------------------------------
# YOKE CONFIG
#------------------------------------------------------------------------------

# these configuration options have been removed from their standard location and
# placed here so that Pagoda Box could override them with the neccessary values
# to configure redundancy.

# IMPORTANT: these settings will always be overriden when the server boots. They
# are set dynamically and so should never change.

listen_addresses = '%s'      # what IP address(es) to listen on;
                                  # comma-separated list of addresses;
                                  # defaults to 'localhost'; use '*' for all
                                  # (change requires restart)
port = %d                     # (change requires restart)
wal_level = hot_standby           # minimal, archive, or hot_standby
                                  # (change requires restart)
archive_mode = on                 # allows archiving to be done
                                  # (change requires restart)
archive_command = 'exit 0'        # command to use to archive a logfile segment
                                  # placeholders: \%p = path of file to archive
                                  #               \%f = file name only
                                  # e.g. 'test ! -f /mnt/server/archivedir/\%f && cp \%p /mnt/server/archivedir/\%f'
max_wal_senders = 10              # max number of walsender processes
                                  # (change requires restart)
wal_keep_segments = 16          	# in logfile segments, 16MB each; 0 disables
hot_standby = on                  # "on" allows queries during recovery
                                  # (change requires restart)
`, opts.listenAddr, conf.PGPort)

	// if this node is currenty 'master' then write one additional configuration
	// into the 'entry'
	if opts.master {
		entry += `
# added only the the node running postgres as 'master'
synchronous_standby_names = slave # standby servers that provide sync rep
                                  # comma-separated list of application_name
                                  # from standby(s); '*' = all
`
	}

	// write 'entry' to the file
	err = ioutil.WriteFile(file, []byte(entry), 0644)
	if err != nil {
		log.Error("[pgconfig.configurePGConf] Failed to write to '%s'!\n%s\n", file, err)
	}

	return nil
}

// createRecovery creates a 'recovery.conf' file with the necessary settings
// required for redundancy on Pagoda Box. This method is called on the node that
// is being configured to run the 'slave' instance of postgres
func createRecovery() error {

	file := conf.DataDir + "recovery.conf"

	self := Whoami()
	other, err := Whoisnot(self.CRole)
	if err != nil {
		log.Fatal("[pgconfig.createRecovery] Unable to find another... Exiting!\n%s\n", err)
		os.Exit(1)
	}

	// open/truncate the recover.conf
	f, err := os.Create(file)
	if err != nil {
		log.Error("[pgconfig.createRecovery] Failed to create '%s'!\n%s\n", file, err)
		return err
	}

	// write manual configuration an 'entry'
	entry := fmt.Sprintf(`#~-----------------------------------------------------------------------------
# YOKE CONFIG
#------------------------------------------------------------------------------

# IMPORTANT: this config file is dynamically generated by Pagoda Box for redundancy
# any changes made here will be overriden.

# When standby_mode is enabled, the PostgreSQL server will work as a standby. It
# tries to connect to the primary according to the connection settings
# primary_conninfo, and receives XLOG records continuously.
standby_mode = on
primary_conninfo = 'host=%s port=%d application_name=slave'

# restore_command specifies the shell command that is executed to copy log files
# back from archival storage. This parameter is *required* for an archive
# recovery, but optional for streaming replication. The given command satisfies
# the requirement without doing anything.
restore_command = 'exit 0'
`, other.Ip, other.PGPort)

	// write 'entry' to the file
	if _, err := f.WriteString(entry); err != nil {
		log.Error("[pgconfig.createRecovery] Failed to write to '%s'!\n%s\n", file, err)
		return err
	}

	return nil
}

// destroyRecovery attempts to destroy the 'recovery.conf'. This method is called
// on a node that is being configured to run the 'master' instance of postgres
func destroyRecovery() {

	file := conf.DataDir + "recovery.conf"

	// remove 'recovery.conf'
	err := os.Remove(file)
	if err != nil {
		log.Warn("[pgconfig.destroyRecovery] No recovery.conf found at '%s'", file)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package handlers

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/security"
	"v.io/v23/vom"
	"v.io/x/ref/services/identity/internal/blesser"
	"v.io/x/ref/services/identity/internal/util"
)

const (
	publicKeyFormKey    = "public_key"
	tokenFormKey        = "token"
	caveatsFormKey      = "caveats"
	outputFormatFormKey = "output_format"

	jsonFormat      = "json"
	base64VomFormat = "base64vom"
)

type accessTokenBlesser struct {
	ctx    *context.T
	params blesser.OAuthBlesserParams
}

// NewOAuthBlessingHandler returns an http.Handler that uses Google OAuth2 Access tokens
// to obtain the username of the requestor and reponds with blessings for that username.
//
// The blessings are namespaced under the ClientID for the access token. In particular,
// the name of the granted blessing is of the form <idp>/<clientID>/<email> where <idp>
// is the name of the default blessings used by the identity provider.
//
// Blessings generated by this service carry a third-party revocation caveat if a
// RevocationManager is specified by the params or they carry an ExpiryCaveat that
// expires after the duration specified by the params.
//
// The handler expects the following request parameters:
// - "public_key": Base64 DER encoded PKIX representation of the client's public key
// - "caveats": Base64 VOM encoded list of caveats [OPTIONAL]
// - "token": Google OAuth2 Access token
// - "output_format": The encoding format for the returned blessings. The following
//   formats are supported:
//     - "json": JSON-encoding of the wire format of Blessings.
//     - "base64vom": Base64 encoding of VOM-encoded Blessings [DEFAULT]
//
// The response consists of blessings encoded in the requested output format.
//
// WARNINGS:
//   - There is no binding between the channel over which the access token
//     was obtained and the channel used to make this request.
//   - There is no "proof of possession of private key" required by the server.
// Thus, if Mallory (attacker) possesses the access token associated with Alice's
// account (victim), she may be able to obtain a blessing with Alice's name on it
// for any public key of her choice.
func NewOAuthBlessingHandler(ctx *context.T, params blesser.OAuthBlesserParams) http.Handler {
	return &accessTokenBlesser{ctx, params}
}

func (a *accessTokenBlesser) blessingCaveats(r *http.Request, p security.Principal) ([]security.Caveat, error) {
	var caveats []security.Caveat
	if base64VomCaveats := r.FormValue(caveatsFormKey); len(base64VomCaveats) != 0 {
		vomCaveats, err := base64.URLEncoding.DecodeString(base64VomCaveats)
		if err != nil {
			return nil, fmt.Errorf("base64.URLEncoding.DecodeString failed: %v", err)
		}

		if err := vom.Decode(vomCaveats, &caveats); err != nil {
			return nil, fmt.Errorf("vom.Decode failed: %v", err)
		}
	}
	// TODO(suharshs, ataly): Should we ensure that we have at least a
	// revocation or expiry caveat?
	if len(caveats) == 0 {
		var (
			cav security.Caveat
			err error
		)
		if a.params.RevocationManager != nil {
			cav, err = a.params.RevocationManager.NewCaveat(p.PublicKey(), a.params.DischargerLocation)
		} else {
			cav, err = security.NewExpiryCaveat(time.Now().Add(a.params.BlessingDuration))
		}
		if err != nil {
			return nil, fmt.Errorf("failed to construct caveats: %v", err)
		}
		caveats = append(caveats, cav)
	}
	return caveats, nil

}

func (a *accessTokenBlesser) remotePublicKey(r *http.Request) (security.PublicKey, error) {
	publicKeyVom, err := base64.URLEncoding.DecodeString(r.FormValue(publicKeyFormKey))
	if err != nil {
		return nil, fmt.Errorf("base64.URLEncoding.DecodeString failed: %v", err)
	}
	return security.UnmarshalPublicKey(publicKeyVom)
}

func (a *accessTokenBlesser) blessingExtension(r *http.Request) (string, error) {
	email, clientID, err := a.params.OAuthProvider.GetEmailAndClientID(r.FormValue(tokenFormKey))
	if err != nil {
		return "", err
	}
	// We use <clientID>/<email> as the extension in order to namespace the blessing under
	// the <clientID>. This has the downside that the blessing cannot be used to act on
	// behalf of the user, i.e., services access controlled to blessings matching"<idp>/<email>"
	// would not authorize this blessing.
	//
	// The alternative is to use the extension <email>/<clientID> however this is risky as it
	// may provide too much authority to the app, especially since we don't have a set of default
	// caveats to apply to the blessing.
	//
	// TODO(ataly, ashankar): Think about changing to the extension <email>/<clientID>.
	return strings.Join([]string{clientID, email}, security.ChainSeparator), nil
}

func (a *accessTokenBlesser) encodeBlessingsJson(b security.Blessings) ([]byte, error) {
	return json.Marshal(security.MarshalBlessings(b))
}

func (a *accessTokenBlesser) encodeBlessingsVom(b security.Blessings) (string, error) {
	bVom, err := vom.Encode(b)
	if err != nil {
		return "", fmt.Errorf("vom.Encode(%v) failed: %v", b, err)
	}
	return base64.URLEncoding.EncodeToString(bVom), nil
}

func (a *accessTokenBlesser) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	remoteKey, err := a.remotePublicKey(r)
	if err != nil {
		a.ctx.Info("Failed to decode public key [%v] for request %#v", err, r)
		util.HTTPServerError(w, fmt.Errorf("failed to decode public key: %v", err))
		return
	}

	p := v23.GetPrincipal(a.ctx)
	with := p.BlessingStore().Default()

	caveats, err := a.blessingCaveats(r, p)
	if err != nil {
		a.ctx.Info("Failed to constuct caveats for blessing [%v] for request %#v", err, r)
		util.HTTPServerError(w, fmt.Errorf("failed to construct caveats for blessing: %v", err))
		return
	}

	extension, err := a.blessingExtension(r)
	if err != nil {
		a.ctx.Info("Failed to process access token [%v] for request %#v", err, r)
		util.HTTPServerError(w, fmt.Errorf("failed to process access token: %v", err))
		return
	}

	blessings, err := p.Bless(remoteKey, with, extension, caveats[0], caveats[1:]...)
	if err != nil {
		a.ctx.Info("Failed to Bless [%v] for request %#v", err, r)
		util.HTTPServerError(w, fmt.Errorf("failed to Bless: %v", err))
		return
	}

	outputFormat := r.FormValue(outputFormatFormKey)
	if len(outputFormat) == 0 {
		outputFormat = base64VomFormat
	}
	switch outputFormat {
	case jsonFormat:
		encodedBlessings, err := a.encodeBlessingsJson(blessings)
		if err != nil {
			a.ctx.Info("Failed to encode blessings [%v] for request %#v", err, r)
			util.HTTPServerError(w, fmt.Errorf("failed to encode blessings in format %v: %v", outputFormat, err))
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.Write(encodedBlessings)
	case base64VomFormat:
		encodedBlessings, err := a.encodeBlessingsVom(blessings)
		if err != nil {
			a.ctx.Info("Failed to encode blessings [%v] for request %#v", err, r)
			util.HTTPServerError(w, fmt.Errorf("failed to encode blessings in format %v: %v", outputFormat, err))
			return
		}
		w.Header().Set("Content-Type", "application/text")
		w.Write([]byte(encodedBlessings))
	default:
		a.ctx.Info("Unrecognized output format [%v] in request %#v", outputFormat, r)
		util.HTTPServerError(w, fmt.Errorf("unrecognized output format [%v] in request. Allowed formats are [%v, %v]", outputFormat, base64VomFormat, jsonFormat))
		return
	}
}
<commit_msg>services/identity: Fix log line formatting.<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package handlers

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/security"
	"v.io/v23/vom"
	"v.io/x/ref/services/identity/internal/blesser"
	"v.io/x/ref/services/identity/internal/util"
)

const (
	publicKeyFormKey    = "public_key"
	tokenFormKey        = "token"
	caveatsFormKey      = "caveats"
	outputFormatFormKey = "output_format"

	jsonFormat      = "json"
	base64VomFormat = "base64vom"
)

type accessTokenBlesser struct {
	ctx    *context.T
	params blesser.OAuthBlesserParams
}

// NewOAuthBlessingHandler returns an http.Handler that uses Google OAuth2 Access tokens
// to obtain the username of the requestor and reponds with blessings for that username.
//
// The blessings are namespaced under the ClientID for the access token. In particular,
// the name of the granted blessing is of the form <idp>/<clientID>/<email> where <idp>
// is the name of the default blessings used by the identity provider.
//
// Blessings generated by this service carry a third-party revocation caveat if a
// RevocationManager is specified by the params or they carry an ExpiryCaveat that
// expires after the duration specified by the params.
//
// The handler expects the following request parameters:
// - "public_key": Base64 DER encoded PKIX representation of the client's public key
// - "caveats": Base64 VOM encoded list of caveats [OPTIONAL]
// - "token": Google OAuth2 Access token
// - "output_format": The encoding format for the returned blessings. The following
//   formats are supported:
//     - "json": JSON-encoding of the wire format of Blessings.
//     - "base64vom": Base64 encoding of VOM-encoded Blessings [DEFAULT]
//
// The response consists of blessings encoded in the requested output format.
//
// WARNINGS:
//   - There is no binding between the channel over which the access token
//     was obtained and the channel used to make this request.
//   - There is no "proof of possession of private key" required by the server.
// Thus, if Mallory (attacker) possesses the access token associated with Alice's
// account (victim), she may be able to obtain a blessing with Alice's name on it
// for any public key of her choice.
func NewOAuthBlessingHandler(ctx *context.T, params blesser.OAuthBlesserParams) http.Handler {
	return &accessTokenBlesser{ctx, params}
}

func (a *accessTokenBlesser) blessingCaveats(r *http.Request, p security.Principal) ([]security.Caveat, error) {
	var caveats []security.Caveat
	if base64VomCaveats := r.FormValue(caveatsFormKey); len(base64VomCaveats) != 0 {
		vomCaveats, err := base64.URLEncoding.DecodeString(base64VomCaveats)
		if err != nil {
			return nil, fmt.Errorf("base64.URLEncoding.DecodeString failed: %v", err)
		}

		if err := vom.Decode(vomCaveats, &caveats); err != nil {
			return nil, fmt.Errorf("vom.Decode failed: %v", err)
		}
	}
	// TODO(suharshs, ataly): Should we ensure that we have at least a
	// revocation or expiry caveat?
	if len(caveats) == 0 {
		var (
			cav security.Caveat
			err error
		)
		if a.params.RevocationManager != nil {
			cav, err = a.params.RevocationManager.NewCaveat(p.PublicKey(), a.params.DischargerLocation)
		} else {
			cav, err = security.NewExpiryCaveat(time.Now().Add(a.params.BlessingDuration))
		}
		if err != nil {
			return nil, fmt.Errorf("failed to construct caveats: %v", err)
		}
		caveats = append(caveats, cav)
	}
	return caveats, nil

}

func (a *accessTokenBlesser) remotePublicKey(r *http.Request) (security.PublicKey, error) {
	publicKeyVom, err := base64.URLEncoding.DecodeString(r.FormValue(publicKeyFormKey))
	if err != nil {
		return nil, fmt.Errorf("base64.URLEncoding.DecodeString failed: %v", err)
	}
	return security.UnmarshalPublicKey(publicKeyVom)
}

func (a *accessTokenBlesser) blessingExtension(r *http.Request) (string, error) {
	email, clientID, err := a.params.OAuthProvider.GetEmailAndClientID(r.FormValue(tokenFormKey))
	if err != nil {
		return "", err
	}
	// We use <clientID>/<email> as the extension in order to namespace the blessing under
	// the <clientID>. This has the downside that the blessing cannot be used to act on
	// behalf of the user, i.e., services access controlled to blessings matching"<idp>/<email>"
	// would not authorize this blessing.
	//
	// The alternative is to use the extension <email>/<clientID> however this is risky as it
	// may provide too much authority to the app, especially since we don't have a set of default
	// caveats to apply to the blessing.
	//
	// TODO(ataly, ashankar): Think about changing to the extension <email>/<clientID>.
	return strings.Join([]string{clientID, email}, security.ChainSeparator), nil
}

func (a *accessTokenBlesser) encodeBlessingsJson(b security.Blessings) ([]byte, error) {
	return json.Marshal(security.MarshalBlessings(b))
}

func (a *accessTokenBlesser) encodeBlessingsVom(b security.Blessings) (string, error) {
	bVom, err := vom.Encode(b)
	if err != nil {
		return "", fmt.Errorf("vom.Encode(%v) failed: %v", b, err)
	}
	return base64.URLEncoding.EncodeToString(bVom), nil
}

func (a *accessTokenBlesser) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	remoteKey, err := a.remotePublicKey(r)
	if err != nil {
		a.ctx.Infof("Failed to decode public key [%v] for request %#v", err, r)
		util.HTTPServerError(w, fmt.Errorf("failed to decode public key: %v", err))
		return
	}

	p := v23.GetPrincipal(a.ctx)
	with := p.BlessingStore().Default()

	caveats, err := a.blessingCaveats(r, p)
	if err != nil {
		a.ctx.Infof("Failed to constuct caveats for blessing [%v] for request %#v", err, r)
		util.HTTPServerError(w, fmt.Errorf("failed to construct caveats for blessing: %v", err))
		return
	}

	extension, err := a.blessingExtension(r)
	if err != nil {
		a.ctx.Infof("Failed to process access token [%v] for request %#v", err, r)
		util.HTTPServerError(w, fmt.Errorf("failed to process access token: %v", err))
		return
	}

	blessings, err := p.Bless(remoteKey, with, extension, caveats[0], caveats[1:]...)
	if err != nil {
		a.ctx.Infof("Failed to Bless [%v] for request %#v", err, r)
		util.HTTPServerError(w, fmt.Errorf("failed to Bless: %v", err))
		return
	}

	outputFormat := r.FormValue(outputFormatFormKey)
	if len(outputFormat) == 0 {
		outputFormat = base64VomFormat
	}
	switch outputFormat {
	case jsonFormat:
		encodedBlessings, err := a.encodeBlessingsJson(blessings)
		if err != nil {
			a.ctx.Infof("Failed to encode blessings [%v] for request %#v", err, r)
			util.HTTPServerError(w, fmt.Errorf("failed to encode blessings in format %v: %v", outputFormat, err))
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.Write(encodedBlessings)
	case base64VomFormat:
		encodedBlessings, err := a.encodeBlessingsVom(blessings)
		if err != nil {
			a.ctx.Infof("Failed to encode blessings [%v] for request %#v", err, r)
			util.HTTPServerError(w, fmt.Errorf("failed to encode blessings in format %v: %v", outputFormat, err))
			return
		}
		w.Header().Set("Content-Type", "application/text")
		w.Write([]byte(encodedBlessings))
	default:
		a.ctx.Infof("Unrecognized output format [%v] in request %#v", outputFormat, r)
		util.HTTPServerError(w, fmt.Errorf("unrecognized output format [%v] in request. Allowed formats are [%v, %v]", outputFormat, base64VomFormat, jsonFormat))
		return
	}
}
<|endoftext|>