<commit_before>package homecloud

import (
	"encoding/json"
	"fmt"
	"regexp"
	"time"

	"git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
	"github.com/davecgh/go-spew/spew"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/logger"
	"github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/go-ninja/rpc/json2"
	"github.com/ninjasphere/redigo/redis"
)

var log = logger.GetLogger("HomeCloud")

var conn *ninja.Connection

var thingModel *ThingModel
var deviceModel *DeviceModel
var roomModel *RoomModel
var driverModel *DriverModel

var locationRegexp = regexp.MustCompile("\\$device\\/([A-F0-9]*)\\/[^\\/]*\\/location")

type incomingLocationUpdate struct {
	Zone *string `json:"zone,omitempty"`
}

type outgoingLocationUpdate struct {
	ID         *string `json:"id"`
	HasChanged bool    `json:"hasChanged"`
}

func Start() {

	redisConn, err := redis.Dial("tcp", ":6379")

	if err != nil {
		log.FatalError(err, "Couldn't connect to redis")
	}

	//defer redisConn.Close()

	conn, err = ninja.Connect("sphere-go-homecloud")

	if err != nil {
		log.FatalError(err, "Failed to connect to mqtt")
	}

	thingModel = NewThingModel(redisConn)
	conn.MustExportService(thingModel, "$home/services/ThingModel", &model.ServiceAnnouncement{
		Schema: "/service/thing-model",
	})

	deviceModel = NewDeviceModel(redisConn)
	conn.MustExportService(deviceModel, "$home/services/DeviceModel", &model.ServiceAnnouncement{
		Schema: "/service/device-model",
	})

	roomModel = NewRoomModel(redisConn)
	conn.MustExportService(roomModel, "$home/services/RoomModel", &model.ServiceAnnouncement{
		Schema: "/service/room-model",
	})

	driverModel = NewDriverModel(redisConn)

	startManagingDrivers()
	startManagingDevices()
	startMonitoringLocations()
}

func startDriver(node string, driverID string, config *string) error {

	var rawConfig json.RawMessage
	if config != nil {
		rawConfig = []byte(*config)
	} else {
		rawConfig = []byte("{}")
	}

	client := conn.GetServiceClient(fmt.Sprintf("$node/%s/driver/%s", node, driverID))
	err := client.Call("start", &rawConfig, nil, 10*time.Second)

	if err != nil {
		jsonError, ok := err.(*json2.Error)
		if ok {
			if jsonError.Code == json2.E_INVALID_REQ {

				err := driverModel.DeleteConfig(driverID)
				if err != nil {
					log.Warningf("Driver %s could not parse its config. Also, we couldn't clear it! errors:%s and %s", driverID, jsonError.Message, err)
				} else {
					log.Warningf("Driver %s could not parse its config, so we cleared it from redis. error:%s", driverID, jsonError.Message)
				}

				return startDriver(node, driverID, nil)
			}
		}
	}

	return err
}

func startManagingDrivers() {

	conn.Subscribe("$node/:node/driver/:driver/event/announce", func(announcement *json.RawMessage, values map[string]string) bool {

		node, driver := values["node"], values["driver"]

		log.Infof("Got driver announcement node:%s driver:%s announcement:%s", node, driver, announcement)

		if announcement == nil {
			log.Warningf("Nil driver announcement from node:%s driver:%s", node, driver)
			return true
		}

		module := &model.Module{}
		err := json.Unmarshal(*announcement, module)

		if announcement == nil {
			log.Warningf("Could not parse announcement from node:%s driver:%s error:%s", node, driver, err)
			return true
		}

		err = driverModel.Create(module)
		if err != nil {
			log.Warningf("Failed to save driver announcement for %s error:%s", driver, err)
		}

		config, err := driverModel.GetConfig(values["driver"])

		if err != nil {
			log.Warningf("Failed to retrieve config for driver %s error:%s", driver, err)
		} else {
			err = startDriver(node, driver, config)
			if err != nil {
				log.Warningf("Failed to start driver: %s error:%s", driver, err)
			}
		}

		return true
	})

	conn.Subscribe("$node/:node/driver/:driver/event/config", func(config *json.RawMessage, values map[string]string) bool {
		log.Infof("Got driver config node:%s driver:%s config:%s", values["node"], values["driver"], *config)

		if config != nil {
			err := driverModel.SetConfig(values["driver"], string(*config))

			if err != nil {
				log.Warningf("Failed to save config for driver: %s error: %s", values["driver"], err)
			}
		} else {
			log.Infof("Nil config recevied from node:%s driver:%s", values["node"], values["driver"])
		}

		return true
	})

}

func startManagingDevices() {

	conn.Subscribe("$device/:id/event/announce", func(announcement *json.RawMessage, values map[string]string) bool {

		id := values["id"]

		log.Infof("Got device announcement device:%s announcement:%s", id, announcement)

		if announcement == nil {
			log.Warningf("Nil driver announcement from device:%s", id)
			return true
		}

		device := &model.Device{}
		err := json.Unmarshal(*announcement, device)

		if announcement == nil {
			log.Warningf("Could not parse announcement from device:%s error:%s", id, err)
			return true
		}

		err = deviceModel.Create(device)
		if err != nil {
			log.Warningf("Failed to save device announcement for device:%s error:%s", id, err)
		}

		return true
	})

	conn.Subscribe("$device/:device/channel/:channel/event/announce", func(announcement *json.RawMessage, values map[string]string) bool {

		deviceID, channelID := values["device"], values["channel"]

		log.Infof("Got channel announcement device:%s channel:%s announcement:%s", deviceID, channelID, announcement)

		if announcement == nil {
			log.Warningf("Nil channel announcement from device:%s channel:%s", deviceID, channelID)
			return true
		}

		channel := &model.Channel{}
		err := json.Unmarshal(*announcement, channel)

		if announcement == nil {
			log.Warningf("Could not parse channel announcement from device:%s channel:%s error:%s", deviceID, channelID, err)
			return true
		}

		err = deviceModel.AddChannel(channel)
		if err != nil {
			log.Warningf("Failed to save channel announcement for device:%s channel:%s error:%s", deviceID, channelID, err)
		}

		return true
	})

}

func startMonitoringLocations() {

	filter, err := mqtt.NewTopicFilter("$device/+/+/location", 0)
	if err != nil {
		log.FatalError(err, "Failed to subscribe to device locations")
	}

	receipt, err := conn.GetMqttClient().StartSubscription(func(_ *mqtt.MqttClient, message mqtt.Message) {

		deviceID := locationRegexp.FindAllStringSubmatch(message.Topic(), -1)[0][1]

		update := &incomingLocationUpdate{}
		err := json.Unmarshal(message.Payload(), update)
		if err != nil {
			log.Errorf("Failed to parse location update %s to %s : %s", message.Payload(), message.Topic(), err)
			return
		}

		thing, err := thingModel.FetchByDeviceId(deviceID)
		if err != nil {
			log.FatalError(err, fmt.Sprintf("Failed to fetch thing by device id %s", deviceID))
		}

		if update.Zone == nil {
			log.Debugf("< Incoming location update: device %s not in a zone", deviceID)
		} else {
			log.Debugf("< Incoming location update: device %s is in zone %s", deviceID, *update.Zone)
		}

		hasChangedZone := true

		if thing == nil {
			log.Debugf("Device %s is not attached to a thing. Ignoring.", deviceID)
		} else {

			if (thing.Location != nil && update.Zone != nil && *thing.Location == *update.Zone) || (thing.Location == nil && update.Zone == nil) {
				// It's already there
				log.Debugf("Thing %s (%s) (Device %s) was already in that zone.", thing.ID, thing.Name, deviceID)
				hasChangedZone = true
			} else {

				log.Debugf("Thing %s (%s) (Device %s) moved from %s to %s", thing.ID, thing.Name, deviceID, thing.Location, update.Zone)

				err = roomModel.MoveThing(thing.Location, update.Zone, thing.ID)

				if err != nil {
					log.FatalError(err, fmt.Sprintf("Failed to update location of thing %s", thing.ID))
				}

				err = thingModel.SetLocation(thing.ID, update.Zone)
				if err != nil {
					log.FatalError(err, fmt.Sprintf("Failed to update location property of thing %s", thing.ID))
				}

				if update.Zone != nil {
					room, err := roomModel.Fetch(*update.Zone)
					if err != nil {
						log.FatalError(err, fmt.Sprintf("Failed to fetch room %s", *update.Zone))
					}

					if room == nil {
						// XXX: TODO: Remove me once the cloud room model is sync'd and locatino service uses it
						log.Infof("Unknown room %s. Advising remote location service to forget it.", *update.Zone)

						pubReceipt := conn.GetMqttClient().Publish(mqtt.QoS(0), "$location/delete", message.Payload())
						<-pubReceipt
					}
				}
			}

			topic := fmt.Sprintf("$device/%s/channel/%s/%s/event/state", deviceID, "location", "location")

			payload, _ := json.Marshal(&outgoingLocationUpdate{
				ID:         update.Zone,
				HasChanged: hasChangedZone,
			})

			pubReceipt := conn.GetMqttClient().Publish(mqtt.QoS(0), topic, payload)
			<-pubReceipt

		}

	}, filter)

	if err != nil {
		log.Fatalf("Failed to subscribe to device locations: %s", err)
	}

	<-receipt
}
<commit_msg>Fix build error<commit_after>package homecloud

import (
	"encoding/json"
	"fmt"
	"regexp"
	"time"

	"git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/logger"
	"github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/go-ninja/rpc/json2"
	"github.com/ninjasphere/redigo/redis"
)

var log = logger.GetLogger("HomeCloud")

var conn *ninja.Connection

var thingModel *ThingModel
var deviceModel *DeviceModel
var roomModel *RoomModel
var driverModel *DriverModel

var locationRegexp = regexp.MustCompile("\\$device\\/([A-F0-9]*)\\/[^\\/]*\\/location")

type incomingLocationUpdate struct {
	Zone *string `json:"zone,omitempty"`
}

type outgoingLocationUpdate struct {
	ID         *string `json:"id"`
	HasChanged bool    `json:"hasChanged"`
}

func Start() {

	redisConn, err := redis.Dial("tcp", ":6379")

	if err != nil {
		log.FatalError(err, "Couldn't connect to redis")
	}

	//defer redisConn.Close()

	conn, err = ninja.Connect("sphere-go-homecloud")

	if err != nil {
		log.FatalError(err, "Failed to connect to mqtt")
	}

	thingModel = NewThingModel(redisConn)
	conn.MustExportService(thingModel, "$home/services/ThingModel", &model.ServiceAnnouncement{
		Schema: "/service/thing-model",
	})

	deviceModel = NewDeviceModel(redisConn)
	conn.MustExportService(deviceModel, "$home/services/DeviceModel", &model.ServiceAnnouncement{
		Schema: "/service/device-model",
	})

	roomModel = NewRoomModel(redisConn)
	conn.MustExportService(roomModel, "$home/services/RoomModel", &model.ServiceAnnouncement{
		Schema: "/service/room-model",
	})

	driverModel = NewDriverModel(redisConn)

	startManagingDrivers()
	startManagingDevices()
	startMonitoringLocations()
}

func startDriver(node string, driverID string, config *string) error {

	var rawConfig json.RawMessage
	if config != nil {
		rawConfig = []byte(*config)
	} else {
		rawConfig = []byte("{}")
	}

	client := conn.GetServiceClient(fmt.Sprintf("$node/%s/driver/%s", node, driverID))
	err := client.Call("start", &rawConfig, nil, 10*time.Second)

	if err != nil {
		jsonError, ok := err.(*json2.Error)
		if ok {
			if jsonError.Code == json2.E_INVALID_REQ {

				err := driverModel.DeleteConfig(driverID)
				if err != nil {
					log.Warningf("Driver %s could not parse its config. Also, we couldn't clear it! errors:%s and %s", driverID, jsonError.Message, err)
				} else {
					log.Warningf("Driver %s could not parse its config, so we cleared it from redis. error:%s", driverID, jsonError.Message)
				}

				return startDriver(node, driverID, nil)
			}
		}
	}

	return err
}

func startManagingDrivers() {

	conn.Subscribe("$node/:node/driver/:driver/event/announce", func(announcement *json.RawMessage, values map[string]string) bool {

		node, driver := values["node"], values["driver"]

		log.Infof("Got driver announcement node:%s driver:%s announcement:%s", node, driver, announcement)

		if announcement == nil {
			log.Warningf("Nil driver announcement from node:%s driver:%s", node, driver)
			return true
		}

		module := &model.Module{}
		err := json.Unmarshal(*announcement, module)

		if announcement == nil {
			log.Warningf("Could not parse announcement from node:%s driver:%s error:%s", node, driver, err)
			return true
		}

		err = driverModel.Create(module)
		if err != nil {
			log.Warningf("Failed to save driver announcement for %s error:%s", driver, err)
		}

		config, err := driverModel.GetConfig(values["driver"])

		if err != nil {
			log.Warningf("Failed to retrieve config for driver %s error:%s", driver, err)
		} else {
			err = startDriver(node, driver, config)
			if err != nil {
				log.Warningf("Failed to start driver: %s error:%s", driver, err)
			}
		}

		return true
	})

	conn.Subscribe("$node/:node/driver/:driver/event/config", func(config *json.RawMessage, values map[string]string) bool {
		log.Infof("Got driver config node:%s driver:%s config:%s", values["node"], values["driver"], *config)

		if config != nil {
			err := driverModel.SetConfig(values["driver"], string(*config))

			if err != nil {
				log.Warningf("Failed to save config for driver: %s error: %s", values["driver"], err)
			}
		} else {
			log.Infof("Nil config recevied from node:%s driver:%s", values["node"], values["driver"])
		}

		return true
	})

}

func startManagingDevices() {

	conn.Subscribe("$device/:id/event/announce", func(announcement *json.RawMessage, values map[string]string) bool {

		id := values["id"]

		log.Infof("Got device announcement device:%s announcement:%s", id, announcement)

		if announcement == nil {
			log.Warningf("Nil driver announcement from device:%s", id)
			return true
		}

		device := &model.Device{}
		err := json.Unmarshal(*announcement, device)

		if announcement == nil {
			log.Warningf("Could not parse announcement from device:%s error:%s", id, err)
			return true
		}

		err = deviceModel.Create(device)
		if err != nil {
			log.Warningf("Failed to save device announcement for device:%s error:%s", id, err)
		}

		return true
	})

	conn.Subscribe("$device/:device/channel/:channel/event/announce", func(announcement *json.RawMessage, values map[string]string) bool {

		deviceID, channelID := values["device"], values["channel"]

		log.Infof("Got channel announcement device:%s channel:%s announcement:%s", deviceID, channelID, announcement)

		if announcement == nil {
			log.Warningf("Nil channel announcement from device:%s channel:%s", deviceID, channelID)
			return true
		}

		channel := &model.Channel{}
		err := json.Unmarshal(*announcement, channel)

		if announcement == nil {
			log.Warningf("Could not parse channel announcement from device:%s channel:%s error:%s", deviceID, channelID, err)
			return true
		}

		err = deviceModel.AddChannel(channel)
		if err != nil {
			log.Warningf("Failed to save channel announcement for device:%s channel:%s error:%s", deviceID, channelID, err)
		}

		return true
	})

}

func startMonitoringLocations() {

	filter, err := mqtt.NewTopicFilter("$device/+/+/location", 0)
	if err != nil {
		log.FatalError(err, "Failed to subscribe to device locations")
	}

	receipt, err := conn.GetMqttClient().StartSubscription(func(_ *mqtt.MqttClient, message mqtt.Message) {

		deviceID := locationRegexp.FindAllStringSubmatch(message.Topic(), -1)[0][1]

		update := &incomingLocationUpdate{}
		err := json.Unmarshal(message.Payload(), update)
		if err != nil {
			log.Errorf("Failed to parse location update %s to %s : %s", message.Payload(), message.Topic(), err)
			return
		}

		thing, err := thingModel.FetchByDeviceId(deviceID)
		if err != nil {
			log.FatalError(err, fmt.Sprintf("Failed to fetch thing by device id %s", deviceID))
		}

		if update.Zone == nil {
			log.Debugf("< Incoming location update: device %s not in a zone", deviceID)
		} else {
			log.Debugf("< Incoming location update: device %s is in zone %s", deviceID, *update.Zone)
		}

		hasChangedZone := true

		if thing == nil {
			log.Debugf("Device %s is not attached to a thing. Ignoring.", deviceID)
		} else {

			if (thing.Location != nil && update.Zone != nil && *thing.Location == *update.Zone) || (thing.Location == nil && update.Zone == nil) {
				// It's already there
				log.Debugf("Thing %s (%s) (Device %s) was already in that zone.", thing.ID, thing.Name, deviceID)
				hasChangedZone = true
			} else {

				log.Debugf("Thing %s (%s) (Device %s) moved from %s to %s", thing.ID, thing.Name, deviceID, thing.Location, update.Zone)

				err = roomModel.MoveThing(thing.Location, update.Zone, thing.ID)

				if err != nil {
					log.FatalError(err, fmt.Sprintf("Failed to update location of thing %s", thing.ID))
				}

				err = thingModel.SetLocation(thing.ID, update.Zone)
				if err != nil {
					log.FatalError(err, fmt.Sprintf("Failed to update location property of thing %s", thing.ID))
				}

				if update.Zone != nil {
					room, err := roomModel.Fetch(*update.Zone)
					if err != nil {
						log.FatalError(err, fmt.Sprintf("Failed to fetch room %s", *update.Zone))
					}

					if room == nil {
						// XXX: TODO: Remove me once the cloud room model is sync'd and locatino service uses it
						log.Infof("Unknown room %s. Advising remote location service to forget it.", *update.Zone)

						pubReceipt := conn.GetMqttClient().Publish(mqtt.QoS(0), "$location/delete", message.Payload())
						<-pubReceipt
					}
				}
			}

			topic := fmt.Sprintf("$device/%s/channel/%s/%s/event/state", deviceID, "location", "location")

			payload, _ := json.Marshal(&outgoingLocationUpdate{
				ID:         update.Zone,
				HasChanged: hasChangedZone,
			})

			pubReceipt := conn.GetMqttClient().Publish(mqtt.QoS(0), topic, payload)
			<-pubReceipt

		}

	}, filter)

	if err != nil {
		log.Fatalf("Failed to subscribe to device locations: %s", err)
	}

	<-receipt
}
<|endoftext|>


<commit_before>package receiver

import (
	"github.com/tgermain/grandRepositorySky/communicator"
	"github.com/tgermain/grandRepositorySky/communicator/sender"
	"github.com/tgermain/grandRepositorySky/node"
	"github.com/tgermain/grandRepositorySky/shared"
	"net"
	"runtime"
	// "time"
)

//Objects parts ---------------------------------------------------------

type ReceiverLink struct {
	node   *node.DHTnode
	sender *sender.SenderLink
}

//Private methods -------------------------------------------------------
func (r *ReceiverLink) handleRequest(payload []byte) {
	//unmarshall message
	msg := communicator.UnmarshallMessage(payload)
	//switch depending of message type
	shared.Logger.Debug("Handle Request receive something : %#v", msg)
	switch {
	case msg.TypeOfMsg == communicator.LOOKUP:
		{
			r.receiveLookup(&msg)
		}
	case msg.TypeOfMsg == communicator.LOOKUPRESPONSE:
		{
			r.receiveLookupResponse(&msg)
		}
	case msg.TypeOfMsg == communicator.JOINRING:
		{
			r.receiveJoinRing(&msg)
		}
	case msg.TypeOfMsg == communicator.UPDATESUCCESSOR:
		{
			r.receiveUpdateSuccessor(&msg)
		}
	case msg.TypeOfMsg == communicator.UPDATEPREDECESSOR:
		{
			r.receiveUpdatePredecessor(&msg)
		}
	case msg.TypeOfMsg == communicator.PRINTRING:
		{
			r.receivePrintRing(&msg)
		}
	case msg.TypeOfMsg == communicator.UPDATEFINGERTABLE:
		{
			r.receiveUpdateFingerTable(&msg)
		}
	case msg.TypeOfMsg == communicator.AREYOUALIVE:
		{
			r.receiveHeartBeat(&msg)
		}
	case msg.TypeOfMsg == communicator.IAMALIVE:
		{
			r.receiveHeartBeatResponse(&msg)
		}
	case msg.TypeOfMsg == communicator.GETSUCCESORE:
		{
			r.receiveGetSuccesor(&msg)
		}
	case msg.TypeOfMsg == communicator.GETSUCCESORERESPONSE:
		{
			r.receiveGetSuccesorResponse(&msg)
		}
	case msg.TypeOfMsg == communicator.GETDATA:
		{
			r.receiveGetData(&msg)
		}
	case msg.TypeOfMsg == communicator.GETDATARESPONSE:
		{
			r.receiveGetDataResponse(&msg)
		}
	default:
		{
			//rejected mesage
			shared.Logger.Error("Message rejected")
		}
	}
	// multiple launch a go routine
}

//========RECEIVE
func (r *ReceiverLink) receiveUpdatePredecessor(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "newNodeID", "newNodeIp", "newNodePort") {
		newNodeID, _ := msg.Parameters["newNodeID"]
		newNodeIp, _ := msg.Parameters["newNodeIp"]
		newNodePort, _ := msg.Parameters["newNodePort"]
		shared.Logger.Info("Receive an update Predecessor to %s", newNodeID)

		r.node.UpdatePredecessor(&shared.DistantNode{
			newNodeID,
			newNodeIp,
			newNodePort,
		})

	}
}

func (r *ReceiverLink) receiveUpdateSuccessor(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "newNodeID", "newNodeIp", "newNodePort") {
		newNodeID, _ := msg.Parameters["newNodeID"]
		newNodeIp, _ := msg.Parameters["newNodeIp"]
		newNodePort, _ := msg.Parameters["newNodePort"]
		shared.Logger.Info("Receive an update Successor %s", newNodeID)

		r.node.UpdateSuccessor(&shared.DistantNode{
			newNodeID,
			newNodeIp,
			newNodePort,
		})

	}
}

func (r *ReceiverLink) receivePrintRing(msg *communicator.Message) {
	//write your info and if the successor is the origine of the communicator.Message, send it back to him
	if checkRequiredParams(msg.Parameters, "currentString") {
		currentString, _ := msg.Parameters["currentString"]

		shared.Logger.Info("Receiving a print ring request from %s", msg.Origin.Id)
		if shared.LocalId == msg.Origin.Id {
			shared.Logger.Info("And %s is me !", msg.Origin.Id)
			//I launch this request know print the result
			shared.Logger.Info("The ring is like :\n%s", currentString)
		} else {
			//pass the request around
			r.node.PrintNodeName(&currentString)
			msg.Parameters["currentString"] = currentString
			go r.sender.RelayPrintRing(r.node.GetSuccesor(), msg)
		}
	}
}

func (r *ReceiverLink) receiveJoinRing(msg *communicator.Message) {
	shared.Logger.Info("Receiving join ring message from %s", msg.Origin)
	go r.node.AddToRing(&msg.Origin)
}

func (r *ReceiverLink) receiveLookup(msg *communicator.Message) {
	//check if the parameter are correct
	if checkRequiredParams(msg.Parameters, "idAnswer", "idSearched") {
		idSearched, _ := msg.Parameters["idSearched"]
		idAnswer, _ := msg.Parameters["idAnswer"]
		shared.Logger.Info("Receive a lookup for : %s", idSearched)

		//Am I responsible for the key requested  ?
		if r.node.IsResponsible(idSearched) {
			shared.Logger.Info("I'm responsible !")
			go r.sender.SendLookupResponse(&msg.Origin, idAnswer, idSearched)
		} else {
			//no -> sending the request to the closest node
			shared.Logger.Info("relay the lookup")
			go r.sender.RelayLookup(r.node.FindClosestNode(idSearched), msg)
		}

	}

}

func (r *ReceiverLink) receiveLookupResponse(msg *communicator.Message) {
	//heck if everything required is here
	if checkRequiredParams(msg.Parameters, "idAnswer", "idSearched") {
		idSearched, _ := msg.Parameters["idSearched"]
		idAnswer, _ := msg.Parameters["idAnswer"]

		shared.Logger.Info("Receive a lookup response for : %s", idSearched)

		chanResp, ok2 := communicator.PendingLookups[idAnswer]
		if ok2 {
			chanResp <- msg.Origin
		}
	}
}

func (r *ReceiverLink) receiveUpdateFingerTable(msg *communicator.Message) {
	shared.Logger.Info("Receiving update finger table from %s", msg.Origin.Id)
	r.node.UpdateFingerTable()
}

func (r *ReceiverLink) receiveHeartBeat(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer") {
		shared.Logger.Info("Receiving a heartBeat from %s", msg.Origin.Id)
		idAnswer, _ := msg.Parameters["idAnswer"]

		go r.sender.SendHeartBeatResponse(&msg.Origin, idAnswer)
	}
}

func (r *ReceiverLink) receiveHeartBeatResponse(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer") {
		idAnswer, _ := msg.Parameters["idAnswer"]
		shared.Logger.Info("Receiving a heartBeat response from %s for %s", msg.Origin.Id, idAnswer)

		chanResp, ok2 := communicator.PendingHearBeat[idAnswer]
		if ok2 {
			chanResp <- msg.Origin
		}
	}
}

func (r *ReceiverLink) receiveGetSuccesor(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer") {
		shared.Logger.Info("Receiving a get successor from %s", msg.Origin.Id)
		idAnswer, _ := msg.Parameters["idAnswer"]

		go r.sender.SendGetSuccResponse(&msg.Origin, idAnswer, r.node.GetSuccesor())
	}
}

func (r *ReceiverLink) receiveGetSuccesorResponse(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer", "succSuccID",
		"succSuccIp",
		"succSuccPort") {
		idAnswer, _ := msg.Parameters["idAnswer"]
		succSuccID, _ := msg.Parameters["succSuccID"]
		succSuccIp, _ := msg.Parameters["succSuccIp"]
		succSuccPort, _ := msg.Parameters["succSuccPort"]

		shared.Logger.Info("Receiving a GetSuccesor response from %s for %s", msg.Origin.Id, idAnswer)

		succSucc := shared.DistantNode{
			succSuccID,
			succSuccIp,
			succSuccPort,
		}

		chanResp, ok2 := communicator.PendingGetSucc[idAnswer]
		if ok2 {
			chanResp <- succSucc
		}
	}
}

func (r *ReceiverLink) receiveGetData(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "keySearched", "idSearched") {
		shared.Logger.Info("Receiving a get data from %s", msg.Origin.Id)
		keySearched, _ := msg.Parameters["keySearched"]
		idSearched, _ := msg.Parameters["idSearched"]
		_, forced := msg.Parameters["forced"]

		var result string
		if forced {
			result = r.node.GetLocalData(idSearched)
		} else {
			result = r.node.GetData(idSearched)
		}
		r.sender.SendGetDataResponse(&msg.Origin, keySearched, result)
	}
}

func (r *ReceiverLink) receiveGetDataResponse(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer", "value") {
		idAnswer, _ := msg.Parameters["idAnswer"]
		value, _ := msg.Parameters["value"]

		shared.Logger.Info("Receiving a get data response from %s for %s", msg.Origin.Id, idAnswer)

		chanResp, ok2 := communicator.PendingGetData[idAnswer]
		if ok2 {
			chanResp <- value
		}
	}
}

func (r *ReceiverLink) receiveSetData(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "key", "value") {
		shared.Logger.Info("Receiving a set data from %s", msg.Origin.Id)
		key, _ := msg.Parameters["key"]
		value, _ := msg.Parameters["value"]
		tag, forced := msg.Parameters["forced"]

		if forced {
			r.node.SetLocalData(key, value, tag)
		} else {
			r.node.SetData(key, value)
		}
	}
}

func checkRequiredParams(params map[string]string, p ...string) bool {
	for _, v := range p {
		_, ok := params[v]
		if !ok {
			shared.Logger.Error("missing parameter %s", v)
			return false
		}
	}
	return true
}

//Exported methods ------------------------------------------------------

func (r *ReceiverLink) StartAndListen() {

	//launch a go routine and start to listen on local address
	//handle incoming communicator.Message

	addr, err := net.ResolveUDPAddr("udp", (shared.LocalIp + ":" + shared.LocalPort))
	if err != nil {
		shared.Logger.Critical("error when resolving udp address:", err)
		panic(err)
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		shared.Logger.Critical("error when connecting to udp:", err)
		panic(err)
	}
	// defer conn.Close()
	go func() {
		shared.Logger.Info("Receiver starting to listen on address [%s]", addr)
		for {
			//multiple goroutine ! work !
			buffer := make([]byte, 1024)
			bytesReads, err := conn.Read(buffer)
			if err != nil {
				shared.Logger.Critical("error while reading:", err)
				panic(err)
			}
			payload := buffer[0:bytesReads]
			go r.handleRequest(payload)
			// time.Sleep(time.Millisecond * 10)
			runtime.Gosched()
		}
	}()

}

func MakeReceiver(n *node.DHTnode, s *sender.SenderLink) *ReceiverLink {
	return &ReceiverLink{
		n,
		s,
	}
}
<commit_msg>forgot to add SETDATA to handled messages<commit_after>package receiver

import (
	"github.com/tgermain/grandRepositorySky/communicator"
	"github.com/tgermain/grandRepositorySky/communicator/sender"
	"github.com/tgermain/grandRepositorySky/node"
	"github.com/tgermain/grandRepositorySky/shared"
	"net"
	"runtime"
	// "time"
)

//Objects parts ---------------------------------------------------------

type ReceiverLink struct {
	node   *node.DHTnode
	sender *sender.SenderLink
}

//Private methods -------------------------------------------------------
func (r *ReceiverLink) handleRequest(payload []byte) {
	//unmarshall message
	msg := communicator.UnmarshallMessage(payload)
	//switch depending of message type
	shared.Logger.Debug("Handle Request receive something : %#v", msg)
	switch {
	case msg.TypeOfMsg == communicator.LOOKUP:
		{
			r.receiveLookup(&msg)
		}
	case msg.TypeOfMsg == communicator.LOOKUPRESPONSE:
		{
			r.receiveLookupResponse(&msg)
		}
	case msg.TypeOfMsg == communicator.JOINRING:
		{
			r.receiveJoinRing(&msg)
		}
	case msg.TypeOfMsg == communicator.UPDATESUCCESSOR:
		{
			r.receiveUpdateSuccessor(&msg)
		}
	case msg.TypeOfMsg == communicator.UPDATEPREDECESSOR:
		{
			r.receiveUpdatePredecessor(&msg)
		}
	case msg.TypeOfMsg == communicator.PRINTRING:
		{
			r.receivePrintRing(&msg)
		}
	case msg.TypeOfMsg == communicator.UPDATEFINGERTABLE:
		{
			r.receiveUpdateFingerTable(&msg)
		}
	case msg.TypeOfMsg == communicator.AREYOUALIVE:
		{
			r.receiveHeartBeat(&msg)
		}
	case msg.TypeOfMsg == communicator.IAMALIVE:
		{
			r.receiveHeartBeatResponse(&msg)
		}
	case msg.TypeOfMsg == communicator.GETSUCCESORE:
		{
			r.receiveGetSuccesor(&msg)
		}
	case msg.TypeOfMsg == communicator.GETSUCCESORERESPONSE:
		{
			r.receiveGetSuccesorResponse(&msg)
		}
	case msg.TypeOfMsg == communicator.GETDATA:
		{
			r.receiveGetData(&msg)
		}
	case msg.TypeOfMsg == communicator.GETDATARESPONSE:
		{
			r.receiveGetDataResponse(&msg)
		}
	case msg.TypeOfMsg == communicator.SETDATA:
		{
			r.receiveSetData(&msg)
		}
	default:
		{
			//rejected mesage
			shared.Logger.Error("Message rejected")
		}
	}
	// multiple launch a go routine
}

//========RECEIVE
func (r *ReceiverLink) receiveUpdatePredecessor(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "newNodeID", "newNodeIp", "newNodePort") {
		newNodeID, _ := msg.Parameters["newNodeID"]
		newNodeIp, _ := msg.Parameters["newNodeIp"]
		newNodePort, _ := msg.Parameters["newNodePort"]
		shared.Logger.Info("Receive an update Predecessor to %s", newNodeID)

		r.node.UpdatePredecessor(&shared.DistantNode{
			newNodeID,
			newNodeIp,
			newNodePort,
		})

	}
}

func (r *ReceiverLink) receiveUpdateSuccessor(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "newNodeID", "newNodeIp", "newNodePort") {
		newNodeID, _ := msg.Parameters["newNodeID"]
		newNodeIp, _ := msg.Parameters["newNodeIp"]
		newNodePort, _ := msg.Parameters["newNodePort"]
		shared.Logger.Info("Receive an update Successor %s", newNodeID)

		r.node.UpdateSuccessor(&shared.DistantNode{
			newNodeID,
			newNodeIp,
			newNodePort,
		})

	}
}

func (r *ReceiverLink) receivePrintRing(msg *communicator.Message) {
	//write your info and if the successor is the origine of the communicator.Message, send it back to him
	if checkRequiredParams(msg.Parameters, "currentString") {
		currentString, _ := msg.Parameters["currentString"]

		shared.Logger.Info("Receiving a print ring request from %s", msg.Origin.Id)
		if shared.LocalId == msg.Origin.Id {
			shared.Logger.Info("And %s is me !", msg.Origin.Id)
			//I launch this request know print the result
			shared.Logger.Info("The ring is like :\n%s", currentString)
		} else {
			//pass the request around
			r.node.PrintNodeName(&currentString)
			msg.Parameters["currentString"] = currentString
			go r.sender.RelayPrintRing(r.node.GetSuccesor(), msg)
		}
	}
}

func (r *ReceiverLink) receiveJoinRing(msg *communicator.Message) {
	shared.Logger.Info("Receiving join ring message from %s", msg.Origin)
	go r.node.AddToRing(&msg.Origin)
}

func (r *ReceiverLink) receiveLookup(msg *communicator.Message) {
	//check if the parameter are correct
	if checkRequiredParams(msg.Parameters, "idAnswer", "idSearched") {
		idSearched, _ := msg.Parameters["idSearched"]
		idAnswer, _ := msg.Parameters["idAnswer"]
		shared.Logger.Info("Receive a lookup for : %s", idSearched)

		//Am I responsible for the key requested  ?
		if r.node.IsResponsible(idSearched) {
			shared.Logger.Info("I'm responsible !")
			go r.sender.SendLookupResponse(&msg.Origin, idAnswer, idSearched)
		} else {
			//no -> sending the request to the closest node
			shared.Logger.Info("relay the lookup")
			go r.sender.RelayLookup(r.node.FindClosestNode(idSearched), msg)
		}

	}

}

func (r *ReceiverLink) receiveLookupResponse(msg *communicator.Message) {
	//heck if everything required is here
	if checkRequiredParams(msg.Parameters, "idAnswer", "idSearched") {
		idSearched, _ := msg.Parameters["idSearched"]
		idAnswer, _ := msg.Parameters["idAnswer"]

		shared.Logger.Info("Receive a lookup response for : %s", idSearched)

		chanResp, ok2 := communicator.PendingLookups[idAnswer]
		if ok2 {
			chanResp <- msg.Origin
		}
	}
}

func (r *ReceiverLink) receiveUpdateFingerTable(msg *communicator.Message) {
	shared.Logger.Info("Receiving update finger table from %s", msg.Origin.Id)
	r.node.UpdateFingerTable()
}

func (r *ReceiverLink) receiveHeartBeat(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer") {
		shared.Logger.Info("Receiving a heartBeat from %s", msg.Origin.Id)
		idAnswer, _ := msg.Parameters["idAnswer"]

		go r.sender.SendHeartBeatResponse(&msg.Origin, idAnswer)
	}
}

func (r *ReceiverLink) receiveHeartBeatResponse(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer") {
		idAnswer, _ := msg.Parameters["idAnswer"]
		shared.Logger.Info("Receiving a heartBeat response from %s for %s", msg.Origin.Id, idAnswer)

		chanResp, ok2 := communicator.PendingHearBeat[idAnswer]
		if ok2 {
			chanResp <- msg.Origin
		}
	}
}

func (r *ReceiverLink) receiveGetSuccesor(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer") {
		shared.Logger.Info("Receiving a get successor from %s", msg.Origin.Id)
		idAnswer, _ := msg.Parameters["idAnswer"]

		go r.sender.SendGetSuccResponse(&msg.Origin, idAnswer, r.node.GetSuccesor())
	}
}

func (r *ReceiverLink) receiveGetSuccesorResponse(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer", "succSuccID",
		"succSuccIp",
		"succSuccPort") {
		idAnswer, _ := msg.Parameters["idAnswer"]
		succSuccID, _ := msg.Parameters["succSuccID"]
		succSuccIp, _ := msg.Parameters["succSuccIp"]
		succSuccPort, _ := msg.Parameters["succSuccPort"]

		shared.Logger.Info("Receiving a GetSuccesor response from %s for %s", msg.Origin.Id, idAnswer)

		succSucc := shared.DistantNode{
			succSuccID,
			succSuccIp,
			succSuccPort,
		}

		chanResp, ok2 := communicator.PendingGetSucc[idAnswer]
		if ok2 {
			chanResp <- succSucc
		}
	}
}

func (r *ReceiverLink) receiveGetData(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "keySearched", "idSearched") {
		shared.Logger.Info("Receiving a get data from %s", msg.Origin.Id)
		keySearched, _ := msg.Parameters["keySearched"]
		idSearched, _ := msg.Parameters["idSearched"]
		_, forced := msg.Parameters["forced"]

		var result string
		if forced {
			result = r.node.GetLocalData(idSearched)
		} else {
			result = r.node.GetData(idSearched)
		}
		r.sender.SendGetDataResponse(&msg.Origin, keySearched, result)
	}
}

func (r *ReceiverLink) receiveGetDataResponse(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "idAnswer", "value") {
		idAnswer, _ := msg.Parameters["idAnswer"]
		value, _ := msg.Parameters["value"]

		shared.Logger.Info("Receiving a get data response from %s for %s", msg.Origin.Id, idAnswer)

		chanResp, ok2 := communicator.PendingGetData[idAnswer]
		if ok2 {
			chanResp <- value
		}
	}
}

func (r *ReceiverLink) receiveSetData(msg *communicator.Message) {
	if checkRequiredParams(msg.Parameters, "key", "value") {
		shared.Logger.Info("Receiving a set data from %s", msg.Origin.Id)
		key, _ := msg.Parameters["key"]
		value, _ := msg.Parameters["value"]
		tag, forced := msg.Parameters["forced"]

		if forced {
			r.node.SetLocalData(key, value, tag)
		} else {
			r.node.SetData(key, value)
		}
	}
}

func checkRequiredParams(params map[string]string, p ...string) bool {
	for _, v := range p {
		_, ok := params[v]
		if !ok {
			shared.Logger.Error("missing parameter %s", v)
			return false
		}
	}
	return true
}

//Exported methods ------------------------------------------------------

func (r *ReceiverLink) StartAndListen() {

	//launch a go routine and start to listen on local address
	//handle incoming communicator.Message

	addr, err := net.ResolveUDPAddr("udp", (shared.LocalIp + ":" + shared.LocalPort))
	if err != nil {
		shared.Logger.Critical("error when resolving udp address:", err)
		panic(err)
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		shared.Logger.Critical("error when connecting to udp:", err)
		panic(err)
	}
	// defer conn.Close()
	go func() {
		shared.Logger.Info("Receiver starting to listen on address [%s]", addr)
		for {
			//multiple goroutine ! work !
			buffer := make([]byte, 1024)
			bytesReads, err := conn.Read(buffer)
			if err != nil {
				shared.Logger.Critical("error while reading:", err)
				panic(err)
			}
			payload := buffer[0:bytesReads]
			go r.handleRequest(payload)
			// time.Sleep(time.Millisecond * 10)
			runtime.Gosched()
		}
	}()

}

func MakeReceiver(n *node.DHTnode, s *sender.SenderLink) *ReceiverLink {
	return &ReceiverLink{
		n,
		s,
	}
}
<|endoftext|>


<commit_before>package paxos

type network interface {
	send(msg message)
}

type proposal struct {
	seq  int
	data []byte
}

func (p proposal) lessThan(a proposal) bool { return p.seq < a.seq }

func (p proposal) isEmpty() bool { return p.seq == 0 }

type fsm struct {
	id    int
	peers []int
	nt    network

	propSeq  int
	promises int
	reported proposal
	accepts  int
	chosen   proposal

	// these two fields need to be persisted to stable storage
	promiseSeq int
	accepted   proposal
}

func (sm *fsm) recv(msg message) {
	switch msg.typ {
	case msgPrepare:
		sm.handlePrepare(msg)
	case msgPromise:
		sm.handlePromise(msg)
	case msgAcceptRequest:
		sm.handleAcceptRequest(msg)
	case msgAccepted:
		sm.handleAccepted(msg)
	}
}

func (sm *fsm) propose(seq int, data []byte) {
	if sm.propSeq >= seq {
		panic("expect higher seq")
	}
	sm.propSeq = seq
	sm.promises = 0
	sm.reported = proposal{data: data}
	sm.bcast(message{
		typ: msgPrepare,
		seq: sm.propSeq,
	})
}

func (sm *fsm) chosenValue() ([]byte, bool) {
	if sm.accepted.isEmpty() {
		return nil, false
	} else {
		return sm.accepted.data, true
	}
}

func (sm *fsm) handlePrepare(msg message) {
	if msg.typ != msgPrepare {
		panic("invalid message type")
	}
	if sm.promiseSeq >= msg.seq {
		return
	}
	sm.promiseSeq = msg.seq
	sm.send(msg.from, message{
		typ:  msgPromise,
		seq:  msg.seq,
		prop: sm.accepted,
	})
}

func (sm *fsm) handlePromise(msg message) {
	if msg.typ != msgPromise {
		panic("invalid message type")
	}
	if sm.propSeq != msg.seq {
		return
	}
	sm.promises++
	if sm.reported.lessThan(msg.prop) {
		sm.reported = msg.prop
	}

	if sm.promises != sm.quorum() {
		return
	}
	sm.reported.seq = sm.propSeq
	sm.bcast(message{
		typ:  msgAcceptRequest,
		seq:  sm.propSeq,
		prop: sm.reported,
	})
}

func (sm *fsm) handleAcceptRequest(msg message) {
	if msg.typ != msgAcceptRequest {
		panic("invalid message type")
	}
	if sm.promiseSeq > msg.seq {
		return
	}
	sm.accepted = msg.prop
	sm.send(msg.from, message{
		typ: msgAccepted,
		seq: msg.seq,
	})
}

func (sm *fsm) handleAccepted(msg message) {
	if msg.typ != msgAccepted {
		panic("invalid message type")
	}
	if sm.propSeq != msg.seq {
		return
	}
	sm.accepts++
	if sm.accepts != sm.quorum() {
		return
	}
	sm.chosen = sm.reported
}

func (sm *fsm) bcast(msg message) {
	for _, to := range sm.peers {
		sm.send(to, msg)
	}
}

func (sm *fsm) send(to int, msg message) {
	msg.from = sm.id
	msg.to = to
	sm.nt.send(msg)
}

func (sm *fsm) quorum() int { return len(sm.peers)/2 + 1 }
<commit_msg>paxos: fix chosenValue()<commit_after>package paxos

type network interface {
	send(msg message)
}

type proposal struct {
	seq  int
	data []byte
}

func (p proposal) lessThan(a proposal) bool { return p.seq < a.seq }

func (p proposal) isEmpty() bool { return p.seq == 0 }

type fsm struct {
	id    int
	peers []int
	nt    network

	propSeq  int
	promises int
	reported proposal
	accepts  int
	chosen   proposal

	// these two fields need to be persisted to stable storage
	promiseSeq int
	accepted   proposal
}

func (sm *fsm) recv(msg message) {
	switch msg.typ {
	case msgPrepare:
		sm.handlePrepare(msg)
	case msgPromise:
		sm.handlePromise(msg)
	case msgAcceptRequest:
		sm.handleAcceptRequest(msg)
	case msgAccepted:
		sm.handleAccepted(msg)
	}
}

func (sm *fsm) propose(seq int, data []byte) {
	if sm.propSeq >= seq {
		panic("expect higher seq")
	}
	sm.propSeq = seq
	sm.promises = 0
	sm.reported = proposal{data: data}
	sm.bcast(message{
		typ: msgPrepare,
		seq: sm.propSeq,
	})
}

func (sm *fsm) chosenValue() ([]byte, bool) {
	if sm.chosen.isEmpty() {
		return nil, false
	} else {
		return sm.chosen.data, true
	}
}

func (sm *fsm) handlePrepare(msg message) {
	if msg.typ != msgPrepare {
		panic("invalid message type")
	}
	if sm.promiseSeq >= msg.seq {
		return
	}
	sm.promiseSeq = msg.seq
	sm.send(msg.from, message{
		typ:  msgPromise,
		seq:  msg.seq,
		prop: sm.accepted,
	})
}

func (sm *fsm) handlePromise(msg message) {
	if msg.typ != msgPromise {
		panic("invalid message type")
	}
	if sm.propSeq != msg.seq {
		return
	}
	sm.promises++
	if sm.reported.lessThan(msg.prop) {
		sm.reported = msg.prop
	}

	if sm.promises != sm.quorum() {
		return
	}
	sm.reported.seq = sm.propSeq
	sm.bcast(message{
		typ:  msgAcceptRequest,
		seq:  sm.propSeq,
		prop: sm.reported,
	})
}

func (sm *fsm) handleAcceptRequest(msg message) {
	if msg.typ != msgAcceptRequest {
		panic("invalid message type")
	}
	if sm.promiseSeq > msg.seq {
		return
	}
	sm.accepted = msg.prop
	sm.send(msg.from, message{
		typ: msgAccepted,
		seq: msg.seq,
	})
}

func (sm *fsm) handleAccepted(msg message) {
	if msg.typ != msgAccepted {
		panic("invalid message type")
	}
	if sm.propSeq != msg.seq {
		return
	}
	sm.accepts++
	if sm.accepts != sm.quorum() {
		return
	}
	sm.chosen = sm.reported
}

func (sm *fsm) bcast(msg message) {
	for _, to := range sm.peers {
		sm.send(to, msg)
	}
}

func (sm *fsm) send(to int, msg message) {
	msg.from = sm.id
	msg.to = to
	sm.nt.send(msg)
}

func (sm *fsm) quorum() int { return len(sm.peers)/2 + 1 }
<|endoftext|>


<commit_before>package main

import "fmt"
import "net/http"
import "net/http/httputil"
import "net/url"
import "time"
import "log"
import "text/template"
import "io/ioutil"
import "io"
import "flag"
import "os"
import "github.com/elazarl/goproxy"

func Log(handler http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        log.Printf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
        handler.ServeHTTP(w, r)
    })
}

func IndexPage(wr io.Writer) {
    template_bytes, err := ioutil.ReadFile("templates/index.html")
    if err != nil {
        log.Fatal(err)
    }
    template_string := string(template_bytes[:])

    template_renderer, err := template.New("index").Parse(template_string)
    if err != nil {
        log.Fatal(err)
    }
    template_renderer.Execute(wr, nil)
}

var record = flag.Bool("record", false, "run a http proxy that will record requests")
var playback = flag.String("playback", "", "playback a recoreded session")
var target = flag.String("target", "", "the target host to record")

func getTimeMilis() int64 {
    return time.Now().UnixNano()%1e6/1e3
}

func main() {
    flag.Parse()
    if (*playback == "" && *record == false) {
        fmt.Println("One of --playback <filename> or --record is required")
        flag.Usage()
        os.Exit(1)
    } else if (*record && *target == "") {
        fmt.Println("--target <url> is required with --record")
        flag.Usage()
        os.Exit(1)
    } else if (*record && *target != "") {
        proxy := goproxy.NewProxyHttpServer()
        proxy.Verbose = true
        client := http.Client{}
        fo, err := os.Create(".requests")
        if err != nil { panic(err) }
        // close fo on exit and check for its returned error
        defer func() {
            if err := fo.Close(); err != nil {
                panic(err)
            }
        }()

        startTimeMilis := getTimeMilis()

        proxy.OnRequest().DoFunc(func(r *http.Request,ctx *goproxy.ProxyCtx)(*http.Request,*http.Response) {
            fmt.Println(r.Host)
            request_url, err := url.Parse(r.RequestURI)
            if err != nil {
                panic(err)
            }

            r.RequestURI = ""
            r.URL        = request_url
            r.URL.Scheme = "https"
            r.Host       = *target
            r.URL.Host   = *target
            r.URL.User   = nil

            var requestBytes []byte
            requestBytes, err = httputil.DumpRequest(r, true)
            currentTimeMilis := getTimeMilis()
            fo.WriteString(fmt.Sprint(currentTimeMilis-startTimeMilis))
            fo.Write([]byte{0xff,0xff})
            fo.Write(requestBytes)
            fo.Write([]byte{0xff,0xff,0xff})
            fo.Sync()

            resp, err := client.Do(r)
            if err != nil {
                panic(err)
            }

            return nil, resp
        })
        http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static/"))))
        fmt.Println("Serving on http://localhost:8080")
        log.Fatal(http.ListenAndServe(":8080", proxy))
    }
}
<commit_msg>fix timing<commit_after>package main

import "fmt"
import "net/http"
import "net/http/httputil"
import "net/url"
import "time"
import "log"
import "text/template"
import "io/ioutil"
import "io"
import "flag"
import "os"
import "github.com/elazarl/goproxy"

func Log(handler http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        log.Printf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
        handler.ServeHTTP(w, r)
    })
}

func IndexPage(wr io.Writer) {
    template_bytes, err := ioutil.ReadFile("templates/index.html")
    if err != nil {
        log.Fatal(err)
    }
    template_string := string(template_bytes[:])

    template_renderer, err := template.New("index").Parse(template_string)
    if err != nil {
        log.Fatal(err)
    }
    template_renderer.Execute(wr, nil)
}

var record = flag.Bool("record", false, "run a http proxy that will record requests")
var playback = flag.String("playback", "", "playback a recoreded session")
var target = flag.String("target", "", "the target host to record")

func getTimeMilis() int64 {
    return time.Now().UnixNano()/1e6
}

func runProxy() {
    proxy := goproxy.NewProxyHttpServer()
    proxy.Verbose = true
    client := http.Client{}
    fo, err := os.Create(".requests")
    if err != nil { panic(err) }
    // close fo on exit and check for its returned error
    defer func() {
        if err := fo.Close(); err != nil {
            panic(err)
        }
    }()

    startTimeMilis := getTimeMilis()

    proxy.OnRequest().DoFunc(func(r *http.Request,ctx *goproxy.ProxyCtx)(*http.Request,*http.Response) {
        fmt.Println(r.Host)
        request_url, err := url.Parse(r.RequestURI)
        if err != nil {
            panic(err)
        }

        r.RequestURI = ""
        r.URL        = request_url
        r.URL.Scheme = "https"
        r.Host       = *target
        r.URL.Host   = *target
        r.URL.User   = nil

        var requestBytes []byte
        requestBytes, err = httputil.DumpRequest(r, true)
        currentTimeMilis := getTimeMilis()
        fmt.Println(currentTimeMilis,startTimeMilis)
        fmt.Println(currentTimeMilis-startTimeMilis)
        fo.WriteString(fmt.Sprint(currentTimeMilis-startTimeMilis))
        fo.WriteString("\nLOLPONIES\n")
        fo.Write(requestBytes)
        fo.WriteString("\nLOLPONIES\n")
        fo.Sync()

        resp, err := client.Do(r)
        if err != nil {
            panic(err)
        }

        return nil, resp
    })
    http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static/"))))
    fmt.Println("Serving on http://localhost:8080")
    log.Fatal(http.ListenAndServe(":8080", proxy))
}

func main() {
    flag.Parse()
    if (*playback == "" && *record == false) {
        fmt.Println("One of --playback <filename> or --record is required")
        flag.Usage()
        os.Exit(1)
    } else if (*record && *target == "") {
        fmt.Println("--target <url> is required with --record")
        flag.Usage()
        os.Exit(1)
    } else if (*record && *target != "") {
        runProxy()
    }
}
<|endoftext|>


<commit_before>package db

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/pem"
	"log"
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/kyprizel/ct_mon/models"
)

type MonDBState struct {
	Id         bson.ObjectId `json:"id,omitempty" bson:"_id"`
	StartIndex int64         `bson:"start_index"`
	Created    time.Time     `bson:"created"`
	Updated    time.Time     `bson:"updated"`
}

type CertInfo struct {
	Id                    bson.ObjectId `json:"id,omitempty" bson:"_id"`
	Index                 int64
	CommonName            string    `bson:"CommonName"`
	Issuer                string    `bson:"Issuer"`
	Serial                string    `bson:"Serial"`
	NotBefore             time.Time `bson:"NotBefore"`
	NotAfter              time.Time `bson:"NotAfter"`
	KeyUsage              int       `bson:"KeyUsage"`
	PublicKeyAlgorithm    int       `bson:"PublicKeyAlgorithm"`
	SignatureAlgorithm    int       `bson:"SignatureAlgorithm"`
	DNSNames              []string  `bson:"DNSNames"`
	EmailAddresses        []string  `bson:"EmailAddresses"`
	OCSPServer            []string  `bson:"OCSPServer"`
	IssuingCertificateURL []string  `bson:"IssuingCertificateURL"`
	PEMCert               string    `bson:"pem"`
	Precert               bool      `bson:"precert"`
	Created               time.Time `bson:"created"`
	SHA256Sum             string    `bson:"sha256_sum"`
}

type MonDB struct {
	uri     string
	session *mgo.Session
}

func Init(uri string) (*MonDB, error) {
	session, err := mgo.Dial(uri)
	if err != nil {
		return nil, err
	}
	m := &MonDB{uri: uri, session: session}
	return m, nil
}

func (m *MonDB) getSession() (*mgo.Session, error) {
	if m.session == nil {
		var err error
		m.session, err = mgo.Dial(m.uri)
		if err != nil {
			return nil, err
		}
	}
	return m.session.Clone(), nil
}

func (m *MonDB) LoadState() (int64, error) {
	session, err := m.getSession()
	if err != nil {
		log.Print("DB connection error\n")
		return 0, err
	}
	col := session.DB("").C("state")

	result := MonDBState{}
	err = col.Find(nil).Sort("-updated").One(&result)
	if err != nil {
		return 0, err
	}

	return result.StartIndex, nil
}

func (m *MonDB) SaveState(StartIndex int64) error {
	session, err := m.getSession()
	if err != nil {
		log.Print("DB connection error\n")
		return err
	}

	col := session.DB("").C("state")
	var needInsert bool
	state := MonDBState{}
	err = col.Find(nil).Sort("-updated").One(&state)
	if err != nil {
		needInsert = true
	}

	if !needInsert {
		qs := bson.M{"_id": state.Id}
		change := bson.M{"$set": bson.M{"start_index": StartIndex, "updated": time.Now().UTC()}}
		return col.Update(qs, change)
	}

	state.Id = bson.NewObjectId()
	state.StartIndex = StartIndex
	state.Created = time.Now().UTC()
	state.Updated = state.Created
	return col.Insert(state)
}

func (m *MonDB) StoreCertDetails(cert *CertInfo) error {
	session, err := m.getSession()
	if err != nil {
		log.Print("DB connection error\n")
		return err
	}

	col := session.DB("").C("certificate_details")
	cert.Id = bson.NewObjectId()
	cert.Created = time.Now().UTC()
	return col.Insert(cert)
}

type CertHandler struct {
	DB *MonDB
}

func (s *CertHandler) HandleEvents(ch chan models.MonEvent) {
	for {
		ev := <-ch
		entry := *ev.LogEntry
		switch ev.Type {
		case models.CT_CERT:
			pemCert := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: entry.X509Cert.Raw})
			hasher := sha256.New()
			hasher.Write(entry.X509Cert.Raw)
			sha := hex.EncodeToString(hasher.Sum(nil))
			c := &CertInfo{Index: entry.Index, CommonName: entry.X509Cert.Subject.CommonName,
				Issuer:    entry.X509Cert.Issuer.CommonName,
				Serial:    entry.X509Cert.SerialNumber.String(),
				NotBefore: entry.X509Cert.NotBefore, NotAfter: entry.X509Cert.NotAfter,
				KeyUsage:           int(entry.X509Cert.KeyUsage),
				PublicKeyAlgorithm: int(entry.X509Cert.PublicKeyAlgorithm),
				SignatureAlgorithm: int(entry.X509Cert.SignatureAlgorithm),
				DNSNames:           entry.X509Cert.DNSNames, EmailAddresses: entry.X509Cert.EmailAddresses,
				OCSPServer:            entry.X509Cert.OCSPServer,
				IssuingCertificateURL: entry.X509Cert.IssuingCertificateURL,
				PEMCert:               string(pemCert), Precert: false, SHA256Sum: sha}
			s.DB.StoreCertDetails(c)
		case models.CT_PRECERT:
			pemCert := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: entry.Precert.TBSCertificate.Raw})
			hasher := sha256.New()
			hasher.Write(entry.Precert.TBSCertificate.Raw)
			sha := hex.EncodeToString(hasher.Sum(nil))
			c := &CertInfo{Index: entry.Index, CommonName: entry.Precert.TBSCertificate.Subject.CommonName,
				Issuer:                entry.Precert.TBSCertificate.Issuer.CommonName,
				Serial:                entry.Precert.TBSCertificate.SerialNumber.String(),
				NotBefore:             entry.Precert.TBSCertificate.NotBefore,
				NotAfter:              entry.Precert.TBSCertificate.NotAfter,
				KeyUsage:              int(entry.Precert.TBSCertificate.KeyUsage),
				PublicKeyAlgorithm:    int(entry.Precert.TBSCertificate.PublicKeyAlgorithm),
				SignatureAlgorithm:    int(entry.Precert.TBSCertificate.SignatureAlgorithm),
				DNSNames:              entry.Precert.TBSCertificate.DNSNames,
				EmailAddresses:        entry.Precert.TBSCertificate.EmailAddresses,
				OCSPServer:            entry.Precert.TBSCertificate.OCSPServer,
				IssuingCertificateURL: entry.Precert.TBSCertificate.IssuingCertificateURL,
				PEMCert:               string(pemCert), Precert: true, SHA256Sum: sha}
			s.DB.StoreCertDetails(c)
		case models.CT_QUIT:
			break
		}
	}
}
<commit_msg>check if cert info for index already exists before saving to db<commit_after>package db

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/pem"
	"log"
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	"github.com/kyprizel/ct_mon/models"
)

type MonDBState struct {
	Id         bson.ObjectId `json:"id,omitempty" bson:"_id"`
	StartIndex int64         `bson:"start_index"`
	Created    time.Time     `bson:"created"`
	Updated    time.Time     `bson:"updated"`
}

type CertInfo struct {
	Id                    bson.ObjectId `json:"id,omitempty" bson:"_id"`
	Index                 int64
	CommonName            string    `bson:"CommonName"`
	Issuer                string    `bson:"Issuer"`
	Serial                string    `bson:"Serial"`
	NotBefore             time.Time `bson:"NotBefore"`
	NotAfter              time.Time `bson:"NotAfter"`
	KeyUsage              int       `bson:"KeyUsage"`
	PublicKeyAlgorithm    int       `bson:"PublicKeyAlgorithm"`
	SignatureAlgorithm    int       `bson:"SignatureAlgorithm"`
	DNSNames              []string  `bson:"DNSNames"`
	EmailAddresses        []string  `bson:"EmailAddresses"`
	OCSPServer            []string  `bson:"OCSPServer"`
	IssuingCertificateURL []string  `bson:"IssuingCertificateURL"`
	PEMCert               string    `bson:"pem"`
	Precert               bool      `bson:"precert"`
	Created               time.Time `bson:"created"`
	SHA256Sum             string    `bson:"sha256_sum"`
}

type MonDB struct {
	uri     string
	session *mgo.Session
}

func Init(uri string) (*MonDB, error) {
	session, err := mgo.Dial(uri)
	if err != nil {
		return nil, err
	}
	m := &MonDB{uri: uri, session: session}
	return m, nil
}

func (m *MonDB) getSession() (*mgo.Session, error) {
	if m.session == nil {
		var err error
		m.session, err = mgo.Dial(m.uri)
		if err != nil {
			return nil, err
		}
	}
	return m.session.Clone(), nil
}

func (m *MonDB) LoadState() (int64, error) {
	session, err := m.getSession()
	if err != nil {
		log.Print("DB connection error\n")
		return 0, err
	}
	col := session.DB("").C("state")

	result := MonDBState{}
	err = col.Find(nil).Sort("-updated").One(&result)
	if err != nil {
		return 0, err
	}

	return result.StartIndex, nil
}

func (m *MonDB) SaveState(StartIndex int64) error {
	session, err := m.getSession()
	if err != nil {
		log.Print("DB connection error\n")
		return err
	}

	col := session.DB("").C("state")
	var needInsert bool
	state := MonDBState{}
	err = col.Find(nil).Sort("-updated").One(&state)
	if err != nil {
		needInsert = true
	}

	if !needInsert {
		qs := bson.M{"_id": state.Id}
		change := bson.M{"$set": bson.M{"start_index": StartIndex, "updated": time.Now().UTC()}}
		return col.Update(qs, change)
	}

	state.Id = bson.NewObjectId()
	state.StartIndex = StartIndex
	state.Created = time.Now().UTC()
	state.Updated = state.Created
	return col.Insert(state)
}

func (m *MonDB) StoreCertDetails(cert *CertInfo) error {
	session, err := m.getSession()
	if err != nil {
		log.Print("DB connection error\n")
		return err
	}

	col := session.DB("").C("certificate_details")
	/* do not store same entry more than once */
	if !col.Has(bson.M{"index": cert.Index}) {
		cert.Id = bson.NewObjectId()
		cert.Created = time.Now().UTC()
		return col.Insert(cert)
	}
	return nil;
}

type CertHandler struct {
	DB *MonDB
}

func (s *CertHandler) HandleEvents(ch chan models.MonEvent) {
	for {
		ev := <-ch
		entry := *ev.LogEntry
		switch ev.Type {
		case models.CT_CERT:
			pemCert := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: entry.X509Cert.Raw})
			hasher := sha256.New()
			hasher.Write(entry.X509Cert.Raw)
			sha := hex.EncodeToString(hasher.Sum(nil))
			c := &CertInfo{Index: entry.Index, CommonName: entry.X509Cert.Subject.CommonName,
				Issuer:    entry.X509Cert.Issuer.CommonName,
				Serial:    entry.X509Cert.SerialNumber.String(),
				NotBefore: entry.X509Cert.NotBefore, NotAfter: entry.X509Cert.NotAfter,
				KeyUsage:           int(entry.X509Cert.KeyUsage),
				PublicKeyAlgorithm: int(entry.X509Cert.PublicKeyAlgorithm),
				SignatureAlgorithm: int(entry.X509Cert.SignatureAlgorithm),
				DNSNames:           entry.X509Cert.DNSNames, EmailAddresses: entry.X509Cert.EmailAddresses,
				OCSPServer:            entry.X509Cert.OCSPServer,
				IssuingCertificateURL: entry.X509Cert.IssuingCertificateURL,
				PEMCert:               string(pemCert), Precert: false, SHA256Sum: sha}
			s.DB.StoreCertDetails(c)
		case models.CT_PRECERT:
			pemCert := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: entry.Precert.TBSCertificate.Raw})
			hasher := sha256.New()
			hasher.Write(entry.Precert.TBSCertificate.Raw)
			sha := hex.EncodeToString(hasher.Sum(nil))
			c := &CertInfo{Index: entry.Index, CommonName: entry.Precert.TBSCertificate.Subject.CommonName,
				Issuer:                entry.Precert.TBSCertificate.Issuer.CommonName,
				Serial:                entry.Precert.TBSCertificate.SerialNumber.String(),
				NotBefore:             entry.Precert.TBSCertificate.NotBefore,
				NotAfter:              entry.Precert.TBSCertificate.NotAfter,
				KeyUsage:              int(entry.Precert.TBSCertificate.KeyUsage),
				PublicKeyAlgorithm:    int(entry.Precert.TBSCertificate.PublicKeyAlgorithm),
				SignatureAlgorithm:    int(entry.Precert.TBSCertificate.SignatureAlgorithm),
				DNSNames:              entry.Precert.TBSCertificate.DNSNames,
				EmailAddresses:        entry.Precert.TBSCertificate.EmailAddresses,
				OCSPServer:            entry.Precert.TBSCertificate.OCSPServer,
				IssuingCertificateURL: entry.Precert.TBSCertificate.IssuingCertificateURL,
				PEMCert:               string(pemCert), Precert: true, SHA256Sum: sha}
			s.DB.StoreCertDetails(c)
		case models.CT_QUIT:
			break
		}
	}
}
<|endoftext|>


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

package middleware

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/go-swagger/go-swagger/httpkit"
	"github.com/go-swagger/go-swagger/httpkit/middleware/untyped"
	"github.com/go-swagger/go-swagger/internal/testing/petstore"
	"github.com/go-swagger/go-swagger/loads"
	"github.com/gorilla/context"
	"github.com/stretchr/testify/assert"
)

type stubOperationHandler struct {
}

func (s *stubOperationHandler) ParameterModel() interface{} {
	return nil
}

func (s *stubOperationHandler) Handle(params interface{}) (interface{}, error) {
	return nil, nil
}

type testBinder struct {
}

func (t *testBinder) BindRequest(r *http.Request, m *MatchedRoute) error {
	return nil
}

func TestContentType_Issue264(t *testing.T) {
	swspec, err := loads.Spec("../../fixtures/bugs/264/swagger.yml")
	if assert.NoError(t, err) {
		api := untyped.NewAPI(swspec)
		api.RegisterConsumer("application/json", httpkit.JSONConsumer())
		api.RegisterProducer("application/json", httpkit.JSONProducer())
		api.RegisterOperation("delete", "/key/{id}", new(stubOperationHandler))

		handler := Serve(swspec, api)
		request, _ := http.NewRequest("DELETE", "/key/1", nil)
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, request)
		assert.Equal(t, 200, recorder.Code)
	}
}

func TestServe(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	handler := Serve(spec, api)

	// serve spec document
	request, _ := http.NewRequest("GET", "http://localhost:8080/swagger.json", nil)
	request.Header.Add("Content-Type", httpkit.JSONMime)
	request.Header.Add("Accept", httpkit.JSONMime)
	recorder := httptest.NewRecorder()

	handler.ServeHTTP(recorder, request)
	assert.Equal(t, 200, recorder.Code)

	request, _ = http.NewRequest("GET", "http://localhost:8080/swagger-ui", nil)
	recorder = httptest.NewRecorder()

	handler.ServeHTTP(recorder, request)
	assert.Equal(t, 404, recorder.Code)
}

func TestContextAuthorize(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := httpkit.JSONRequest("GET", "/pets", nil)

	v, ok := context.GetOk(request, ctxSecurityPrincipal)
	assert.False(t, ok)
	assert.Nil(t, v)

	ri, ok := ctx.RouteInfo(request)
	assert.True(t, ok)
	p, err := ctx.Authorize(request, ri)
	assert.Error(t, err)
	assert.Nil(t, p)

	v, ok = context.GetOk(request, ctxSecurityPrincipal)
	assert.False(t, ok)
	assert.Nil(t, v)

	request.SetBasicAuth("wrong", "wrong")
	p, err = ctx.Authorize(request, ri)
	assert.Error(t, err)
	assert.Nil(t, p)

	v, ok = context.GetOk(request, ctxSecurityPrincipal)
	assert.False(t, ok)
	assert.Nil(t, v)

	request.SetBasicAuth("admin", "admin")
	p, err = ctx.Authorize(request, ri)
	assert.NoError(t, err)
	assert.Equal(t, "admin", p)

	v, ok = context.GetOk(request, ctxSecurityPrincipal)
	assert.True(t, ok)
	assert.Equal(t, "admin", v)

	request.SetBasicAuth("doesn't matter", "doesn't")
	pp, rr := ctx.Authorize(request, ri)
	assert.Equal(t, p, pp)
	assert.Equal(t, err, rr)
}

func TestContextNegotiateContentType(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("POST", "/pets", nil)
	// request.Header.Add("Accept", "*/*")
	request.Header.Add("content-type", "text/html")

	v, ok := context.GetOk(request, ctxBoundParams)
	assert.False(t, ok)
	assert.Nil(t, v)

	ri, _ := ctx.RouteInfo(request)

	res := NegotiateContentType(request, ri.Produces, "")
	assert.Equal(t, "", res)

	res2 := NegotiateContentType(request, ri.Produces, "text/plain")
	assert.Equal(t, "text/plain", res2)
}

func TestContextBindAndValidate(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("POST", "/pets", nil)
	request.Header.Add("Accept", "*/*")
	request.Header.Add("content-type", "text/html")
	request.ContentLength = 1

	v, ok := context.GetOk(request, ctxBoundParams)
	assert.False(t, ok)
	assert.Nil(t, v)

	ri, _ := ctx.RouteInfo(request)
	data, result := ctx.BindAndValidate(request, ri) // this requires a much more thorough test
	assert.NotNil(t, data)
	assert.NotNil(t, result)

	v, ok = context.GetOk(request, ctxBoundParams)
	assert.True(t, ok)
	assert.NotNil(t, v)

	dd, rr := ctx.BindAndValidate(request, ri)
	assert.Equal(t, data, dd)
	assert.Equal(t, result, rr)
}

func TestContextRender(t *testing.T) {
	ct := httpkit.JSONMime
	spec, api := petstore.NewAPI(t)

	assert.NotNil(t, spec)
	assert.NotNil(t, api)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "pets", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)
	ri, _ := ctx.RouteInfo(request)

	recorder := httptest.NewRecorder()
	ctx.Respond(recorder, request, []string{ct}, ri, map[string]interface{}{"name": "hello"})
	assert.Equal(t, 200, recorder.Code)
	assert.Equal(t, "{\"name\":\"hello\"}\n", recorder.Body.String())

	recorder = httptest.NewRecorder()
	ctx.Respond(recorder, request, []string{ct}, ri, errors.New("this went wrong"))
	assert.Equal(t, 500, recorder.Code)

	recorder = httptest.NewRecorder()
	assert.Panics(t, func() { ctx.Respond(recorder, request, []string{ct}, ri, map[int]interface{}{1: "hello"}) })

	recorder = httptest.NewRecorder()
	request, _ = http.NewRequest("GET", "pets", nil)
	assert.Panics(t, func() { ctx.Respond(recorder, request, []string{}, ri, map[string]interface{}{"name": "hello"}) })

	request, _ = http.NewRequest("GET", "/pets", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)
	ri, _ = ctx.RouteInfo(request)

	recorder = httptest.NewRecorder()
	ctx.Respond(recorder, request, []string{ct}, ri, map[string]interface{}{"name": "hello"})
	assert.Equal(t, 200, recorder.Code)
	assert.Equal(t, "{\"name\":\"hello\"}\n", recorder.Body.String())

	recorder = httptest.NewRecorder()
	ctx.Respond(recorder, request, []string{ct}, ri, errors.New("this went wrong"))
	assert.Equal(t, 500, recorder.Code)

	recorder = httptest.NewRecorder()
	assert.Panics(t, func() { ctx.Respond(recorder, request, []string{ct}, ri, map[int]interface{}{1: "hello"}) })

	// recorder = httptest.NewRecorder()
	// request, _ = http.NewRequest("GET", "/pets", nil)
	// assert.Panics(t, func() { ctx.Respond(recorder, request, []string{}, ri, map[string]interface{}{"name": "hello"}) })

	recorder = httptest.NewRecorder()
	request, _ = http.NewRequest("DELETE", "/pets/1", nil)
	ri, _ = ctx.RouteInfo(request)
	ctx.Respond(recorder, request, ri.Produces, ri, nil)
	assert.Equal(t, 204, recorder.Code)
}

func TestContextValidResponseFormat(t *testing.T) {
	ct := "application/json"
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)

	// check there's nothing there
	cached, ok := context.GetOk(request, ctxResponseFormat)
	assert.False(t, ok)
	assert.Empty(t, cached)

	// trigger the parse
	mt := ctx.ResponseFormat(request, []string{ct})
	assert.Equal(t, ct, mt)

	// check it was cached
	cached, ok = context.GetOk(request, ctxResponseFormat)
	assert.True(t, ok)
	assert.Equal(t, ct, cached)

	// check if the cast works and fetch from cache too
	mt = ctx.ResponseFormat(request, []string{ct})
	assert.Equal(t, ct, mt)
}

func TestContextInvalidResponseFormat(t *testing.T) {
	ct := "application/x-yaml"
	other := "application/sgml"
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)

	// check there's nothing there
	cached, ok := context.GetOk(request, ctxResponseFormat)
	assert.False(t, ok)
	assert.Empty(t, cached)

	// trigger the parse
	mt := ctx.ResponseFormat(request, []string{other})
	assert.Empty(t, mt)

	// check it was cached
	cached, ok = context.GetOk(request, ctxResponseFormat)
	assert.False(t, ok)
	assert.Empty(t, cached)

	// check if the cast works and fetch from cache too
	mt = ctx.ResponseFormat(request, []string{other})
	assert.Empty(t, mt)
}

func TestContextValidRoute(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "/pets", nil)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxMatchedRoute)
	assert.False(t, ok)

	matched, ok := ctx.RouteInfo(request)
	assert.True(t, ok)
	assert.NotNil(t, matched)

	// check it was cached
	_, ok = context.GetOk(request, ctxMatchedRoute)
	assert.True(t, ok)

	matched, ok = ctx.RouteInfo(request)
	assert.True(t, ok)
	assert.NotNil(t, matched)
}

func TestContextInvalidRoute(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("DELETE", "pets", nil)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxMatchedRoute)
	assert.False(t, ok)

	matched, ok := ctx.RouteInfo(request)
	assert.False(t, ok)
	assert.Nil(t, matched)

	// check it was cached
	_, ok = context.GetOk(request, ctxMatchedRoute)
	assert.False(t, ok)

	matched, ok = ctx.RouteInfo(request)
	assert.False(t, ok)
	assert.Nil(t, matched)
}

func TestContextValidContentType(t *testing.T) {
	ct := "application/json"
	ctx := NewContext(nil, nil, nil)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderContentType, ct)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxContentType)
	assert.False(t, ok)

	// trigger the parse
	mt, _, err := ctx.ContentType(request)
	assert.NoError(t, err)
	assert.Equal(t, ct, mt)

	// check it was cached
	_, ok = context.GetOk(request, ctxContentType)
	assert.True(t, ok)

	// check if the cast works and fetch from cache too
	mt, _, err = ctx.ContentType(request)
	assert.NoError(t, err)
	assert.Equal(t, ct, mt)
}

func TestContextInvalidContentType(t *testing.T) {
	ct := "application("
	ctx := NewContext(nil, nil, nil)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderContentType, ct)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxContentType)
	assert.False(t, ok)

	// trigger the parse
	mt, _, err := ctx.ContentType(request)
	assert.Error(t, err)
	assert.Empty(t, mt)

	// check it was not cached
	_, ok = context.GetOk(request, ctxContentType)
	assert.False(t, ok)

	// check if the failure continues
	_, _, err = ctx.ContentType(request)
	assert.Error(t, err)
}
<commit_msg>moves yaml spec loading functions to its own package<commit_after>// Copyright 2015 go-swagger maintainers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package middleware

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/go-swagger/go-swagger/httpkit"
	"github.com/go-swagger/go-swagger/httpkit/middleware/untyped"
	"github.com/go-swagger/go-swagger/internal/testing/petstore"
	"github.com/go-swagger/go-swagger/loads"
	"github.com/go-swagger/go-swagger/loads/fmts"
	"github.com/gorilla/context"
	"github.com/stretchr/testify/assert"
)

type stubOperationHandler struct {
}

func (s *stubOperationHandler) ParameterModel() interface{} {
	return nil
}

func (s *stubOperationHandler) Handle(params interface{}) (interface{}, error) {
	return nil, nil
}

type testBinder struct {
}

func (t *testBinder) BindRequest(r *http.Request, m *MatchedRoute) error {
	return nil
}

func init() {
	loads.AddLoader(fmts.YAMLMatcher, fmts.YAMLDoc)
}

func TestContentType_Issue264(t *testing.T) {
	swspec, err := loads.Spec("../../fixtures/bugs/264/swagger.yml")
	if assert.NoError(t, err) {
		api := untyped.NewAPI(swspec)
		api.RegisterConsumer("application/json", httpkit.JSONConsumer())
		api.RegisterProducer("application/json", httpkit.JSONProducer())
		api.RegisterOperation("delete", "/key/{id}", new(stubOperationHandler))

		handler := Serve(swspec, api)
		request, _ := http.NewRequest("DELETE", "/key/1", nil)
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, request)
		assert.Equal(t, 200, recorder.Code)
	}
}

func TestServe(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	handler := Serve(spec, api)

	// serve spec document
	request, _ := http.NewRequest("GET", "http://localhost:8080/swagger.json", nil)
	request.Header.Add("Content-Type", httpkit.JSONMime)
	request.Header.Add("Accept", httpkit.JSONMime)
	recorder := httptest.NewRecorder()

	handler.ServeHTTP(recorder, request)
	assert.Equal(t, 200, recorder.Code)

	request, _ = http.NewRequest("GET", "http://localhost:8080/swagger-ui", nil)
	recorder = httptest.NewRecorder()

	handler.ServeHTTP(recorder, request)
	assert.Equal(t, 404, recorder.Code)
}

func TestContextAuthorize(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := httpkit.JSONRequest("GET", "/pets", nil)

	v, ok := context.GetOk(request, ctxSecurityPrincipal)
	assert.False(t, ok)
	assert.Nil(t, v)

	ri, ok := ctx.RouteInfo(request)
	assert.True(t, ok)
	p, err := ctx.Authorize(request, ri)
	assert.Error(t, err)
	assert.Nil(t, p)

	v, ok = context.GetOk(request, ctxSecurityPrincipal)
	assert.False(t, ok)
	assert.Nil(t, v)

	request.SetBasicAuth("wrong", "wrong")
	p, err = ctx.Authorize(request, ri)
	assert.Error(t, err)
	assert.Nil(t, p)

	v, ok = context.GetOk(request, ctxSecurityPrincipal)
	assert.False(t, ok)
	assert.Nil(t, v)

	request.SetBasicAuth("admin", "admin")
	p, err = ctx.Authorize(request, ri)
	assert.NoError(t, err)
	assert.Equal(t, "admin", p)

	v, ok = context.GetOk(request, ctxSecurityPrincipal)
	assert.True(t, ok)
	assert.Equal(t, "admin", v)

	request.SetBasicAuth("doesn't matter", "doesn't")
	pp, rr := ctx.Authorize(request, ri)
	assert.Equal(t, p, pp)
	assert.Equal(t, err, rr)
}

func TestContextNegotiateContentType(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("POST", "/pets", nil)
	// request.Header.Add("Accept", "*/*")
	request.Header.Add("content-type", "text/html")

	v, ok := context.GetOk(request, ctxBoundParams)
	assert.False(t, ok)
	assert.Nil(t, v)

	ri, _ := ctx.RouteInfo(request)

	res := NegotiateContentType(request, ri.Produces, "")
	assert.Equal(t, "", res)

	res2 := NegotiateContentType(request, ri.Produces, "text/plain")
	assert.Equal(t, "text/plain", res2)
}

func TestContextBindAndValidate(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("POST", "/pets", nil)
	request.Header.Add("Accept", "*/*")
	request.Header.Add("content-type", "text/html")
	request.ContentLength = 1

	v, ok := context.GetOk(request, ctxBoundParams)
	assert.False(t, ok)
	assert.Nil(t, v)

	ri, _ := ctx.RouteInfo(request)
	data, result := ctx.BindAndValidate(request, ri) // this requires a much more thorough test
	assert.NotNil(t, data)
	assert.NotNil(t, result)

	v, ok = context.GetOk(request, ctxBoundParams)
	assert.True(t, ok)
	assert.NotNil(t, v)

	dd, rr := ctx.BindAndValidate(request, ri)
	assert.Equal(t, data, dd)
	assert.Equal(t, result, rr)
}

func TestContextRender(t *testing.T) {
	ct := httpkit.JSONMime
	spec, api := petstore.NewAPI(t)

	assert.NotNil(t, spec)
	assert.NotNil(t, api)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "pets", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)
	ri, _ := ctx.RouteInfo(request)

	recorder := httptest.NewRecorder()
	ctx.Respond(recorder, request, []string{ct}, ri, map[string]interface{}{"name": "hello"})
	assert.Equal(t, 200, recorder.Code)
	assert.Equal(t, "{\"name\":\"hello\"}\n", recorder.Body.String())

	recorder = httptest.NewRecorder()
	ctx.Respond(recorder, request, []string{ct}, ri, errors.New("this went wrong"))
	assert.Equal(t, 500, recorder.Code)

	recorder = httptest.NewRecorder()
	assert.Panics(t, func() { ctx.Respond(recorder, request, []string{ct}, ri, map[int]interface{}{1: "hello"}) })

	recorder = httptest.NewRecorder()
	request, _ = http.NewRequest("GET", "pets", nil)
	assert.Panics(t, func() { ctx.Respond(recorder, request, []string{}, ri, map[string]interface{}{"name": "hello"}) })

	request, _ = http.NewRequest("GET", "/pets", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)
	ri, _ = ctx.RouteInfo(request)

	recorder = httptest.NewRecorder()
	ctx.Respond(recorder, request, []string{ct}, ri, map[string]interface{}{"name": "hello"})
	assert.Equal(t, 200, recorder.Code)
	assert.Equal(t, "{\"name\":\"hello\"}\n", recorder.Body.String())

	recorder = httptest.NewRecorder()
	ctx.Respond(recorder, request, []string{ct}, ri, errors.New("this went wrong"))
	assert.Equal(t, 500, recorder.Code)

	recorder = httptest.NewRecorder()
	assert.Panics(t, func() { ctx.Respond(recorder, request, []string{ct}, ri, map[int]interface{}{1: "hello"}) })

	// recorder = httptest.NewRecorder()
	// request, _ = http.NewRequest("GET", "/pets", nil)
	// assert.Panics(t, func() { ctx.Respond(recorder, request, []string{}, ri, map[string]interface{}{"name": "hello"}) })

	recorder = httptest.NewRecorder()
	request, _ = http.NewRequest("DELETE", "/pets/1", nil)
	ri, _ = ctx.RouteInfo(request)
	ctx.Respond(recorder, request, ri.Produces, ri, nil)
	assert.Equal(t, 204, recorder.Code)
}

func TestContextValidResponseFormat(t *testing.T) {
	ct := "application/json"
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)

	// check there's nothing there
	cached, ok := context.GetOk(request, ctxResponseFormat)
	assert.False(t, ok)
	assert.Empty(t, cached)

	// trigger the parse
	mt := ctx.ResponseFormat(request, []string{ct})
	assert.Equal(t, ct, mt)

	// check it was cached
	cached, ok = context.GetOk(request, ctxResponseFormat)
	assert.True(t, ok)
	assert.Equal(t, ct, cached)

	// check if the cast works and fetch from cache too
	mt = ctx.ResponseFormat(request, []string{ct})
	assert.Equal(t, ct, mt)
}

func TestContextInvalidResponseFormat(t *testing.T) {
	ct := "application/x-yaml"
	other := "application/sgml"
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderAccept, ct)

	// check there's nothing there
	cached, ok := context.GetOk(request, ctxResponseFormat)
	assert.False(t, ok)
	assert.Empty(t, cached)

	// trigger the parse
	mt := ctx.ResponseFormat(request, []string{other})
	assert.Empty(t, mt)

	// check it was cached
	cached, ok = context.GetOk(request, ctxResponseFormat)
	assert.False(t, ok)
	assert.Empty(t, cached)

	// check if the cast works and fetch from cache too
	mt = ctx.ResponseFormat(request, []string{other})
	assert.Empty(t, mt)
}

func TestContextValidRoute(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("GET", "/pets", nil)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxMatchedRoute)
	assert.False(t, ok)

	matched, ok := ctx.RouteInfo(request)
	assert.True(t, ok)
	assert.NotNil(t, matched)

	// check it was cached
	_, ok = context.GetOk(request, ctxMatchedRoute)
	assert.True(t, ok)

	matched, ok = ctx.RouteInfo(request)
	assert.True(t, ok)
	assert.NotNil(t, matched)
}

func TestContextInvalidRoute(t *testing.T) {
	spec, api := petstore.NewAPI(t)
	ctx := NewContext(spec, api, nil)
	ctx.router = DefaultRouter(spec, ctx.api)

	request, _ := http.NewRequest("DELETE", "pets", nil)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxMatchedRoute)
	assert.False(t, ok)

	matched, ok := ctx.RouteInfo(request)
	assert.False(t, ok)
	assert.Nil(t, matched)

	// check it was cached
	_, ok = context.GetOk(request, ctxMatchedRoute)
	assert.False(t, ok)

	matched, ok = ctx.RouteInfo(request)
	assert.False(t, ok)
	assert.Nil(t, matched)
}

func TestContextValidContentType(t *testing.T) {
	ct := "application/json"
	ctx := NewContext(nil, nil, nil)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderContentType, ct)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxContentType)
	assert.False(t, ok)

	// trigger the parse
	mt, _, err := ctx.ContentType(request)
	assert.NoError(t, err)
	assert.Equal(t, ct, mt)

	// check it was cached
	_, ok = context.GetOk(request, ctxContentType)
	assert.True(t, ok)

	// check if the cast works and fetch from cache too
	mt, _, err = ctx.ContentType(request)
	assert.NoError(t, err)
	assert.Equal(t, ct, mt)
}

func TestContextInvalidContentType(t *testing.T) {
	ct := "application("
	ctx := NewContext(nil, nil, nil)

	request, _ := http.NewRequest("GET", "http://localhost:8080", nil)
	request.Header.Set(httpkit.HeaderContentType, ct)

	// check there's nothing there
	_, ok := context.GetOk(request, ctxContentType)
	assert.False(t, ok)

	// trigger the parse
	mt, _, err := ctx.ContentType(request)
	assert.Error(t, err)
	assert.Empty(t, mt)

	// check it was not cached
	_, ok = context.GetOk(request, ctxContentType)
	assert.False(t, ok)

	// check if the failure continues
	_, _, err = ctx.ContentType(request)
	assert.Error(t, err)
}
<|endoftext|>


<commit_before>package middlewares

import (
	"github.com/nuveo/prest/config"
	"github.com/nuveo/prest/middlewares"
	"github.com/urfave/negroni"
)

var (
	app *negroni.Negroni

	// MiddlewareStack on pREST
	MiddlewareStack []negroni.Handler

	// BaseStack Middlewares
	BaseStack = []negroni.Handler{
		negroni.Handler(negroni.NewRecovery()),
		negroni.Handler(negroni.NewRecovery()),
		negroni.Handler(middlewares.AccessControl()),
		negroni.Handler(middlewares.HandlerSet()),
	}
)

func initApp() {
	config.InitConf()
	if len(MiddlewareStack) == 0 {
		MiddlewareStack = append(MiddlewareStack, BaseStack...)
	}
	if config.PrestConf.JWTKey != "" {
		MiddlewareStack = append(MiddlewareStack, negroni.Handler(middlewares.JwtMiddleware(config.PrestConf.JWTKey)))
	}
	app = negroni.New(MiddlewareStack...)
}

// GetApp get negroni
func GetApp() *negroni.Negroni {
	if app == nil {
		initApp()
	}
	return app
}
<commit_msg>remove 1 recovery and add logger (#121)<commit_after>package middlewares

import (
	"github.com/nuveo/prest/config"
	"github.com/nuveo/prest/middlewares"
	"github.com/urfave/negroni"
)

var (
	app *negroni.Negroni

	// MiddlewareStack on pREST
	MiddlewareStack []negroni.Handler

	// BaseStack Middlewares
	BaseStack = []negroni.Handler{
		negroni.Handler(negroni.NewRecovery()),
		negroni.Handler(negroni.NewLogger()),
		negroni.Handler(middlewares.AccessControl()),
		negroni.Handler(middlewares.HandlerSet()),
	}
)

func initApp() {
	config.InitConf()
	if len(MiddlewareStack) == 0 {
		MiddlewareStack = append(MiddlewareStack, BaseStack...)
	}
	if config.PrestConf.JWTKey != "" {
		MiddlewareStack = append(MiddlewareStack, negroni.Handler(middlewares.JwtMiddleware(config.PrestConf.JWTKey)))
	}
	app = negroni.New(MiddlewareStack...)
}

// GetApp get negroni
func GetApp() *negroni.Negroni {
	if app == nil {
		initApp()
	}
	return app
}
<|endoftext|>


<commit_before>package windows_test

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/masterzen/winrm"

	. "github.com/onsi/gomega"
)

const dataDir = `C:\var\vcap\data\`

type WindowsEnvironment struct {
	Client *winrm.Client
}

func (e *WindowsEnvironment) ShrinkRootPartition() {
	retryableFailure := "Resize-Partition : Size Not Supported"
	retryableError := "net/http: timeout awaiting response headers"

	cmd := "Get-Partition -DriveLetter C | Resize-Partition -Size $(Get-PartitionSupportedSize -DriveLetter C).SizeMin"

	for i := 0; i < 5; i++ {
		stdout, stderr, exitCode, err := e.RunPowershellCommandWithOffsetAndResponses(
			1,
			cmd,
		)

		if err != nil {
			if strings.Contains(err.Error(), retryableError) {
				fmt.Sprintf("WinRM timed out on attempt %d of 5, waiting 5 seconds to retry command: %s\n", i+1, cmd)
				time.Sleep(5 * time.Second)
				continue
			} else {
				ExpectWithOffset(1, err).NotTo(
					HaveOccurred(),
					fmt.Sprintf(`Command "%s" failed with stdout: %s; stderr: %s`, cmd, stdout, stderr),
				)
			}
		}

		if exitCode != 0 {
			if strings.Contains(stderr, retryableFailure) {
				fmt.Sprintf("Failed to shrink disk on attempt %d of 5, waiting 5 seconds to retry\n", i+1)
				time.Sleep(5 * time.Second)
				continue
			} else {
				ExpectWithOffset(1, exitCode).To(
					BeZero(),
					fmt.Sprintf(
						`Command "%s" failed with exit code: %d; stdout: %s; stderr: %s`, cmd, exitCode, stdout, stderr,
					),
				)
			}
		}

		break
	}
}

func (e *WindowsEnvironment) EnsureRootPartitionAtMaxSize() {
	freeSpaceOutput := e.RunPowershellCommandWithOffset(
		1,
		"Get-Disk $(Get-Partition -DriveLetter C | Select -ExpandProperty DiskNumber) | Select -ExpandProperty LargestFreeExtent",
	)

	freeSpace, err := strconv.Atoi(strings.TrimSpace(freeSpaceOutput))
	ExpectWithOffset(1, err).NotTo(HaveOccurred())

	if freeSpace > 0 {
		e.RunPowershellCommandWithOffset(
			1,
			"Resize-Partition -DriveLetter C -Size $(Get-PartitionSupportedSize -DriveLetter C).SizeMax",
		)
	}
}

func (e *WindowsEnvironment) GetDataDirPartitionNumber() string {
	return strings.TrimSpace(e.RunPowershellCommandWithOffset(
		1,
		fmt.Sprintf(`Get-Partition | Where AccessPaths -Contains "%s" | Select -ExpandProperty PartitionNumber`, dataDir),
	))
}

func (e *WindowsEnvironment) PartitionWithDataDirExists(diskNumber string) bool {
	return e.PartitionWithDataDirExistsWithOffset(1, diskNumber)
}

func (e *WindowsEnvironment) PartitionWithDataDirExistsWithOffset(offset int, diskNumber string) bool {
	stdout := e.RunPowershellCommandWithOffset(
		offset+1,
		fmt.Sprintf(
			`Get-Partition | where AccessPaths -Contains "%s" | Select -ExpandProperty DiskNumber`,
			dataDir,
		),
	)

	return strings.TrimSpace(stdout) == diskNumber
}

func (e *WindowsEnvironment) PartitionWithDataDirExistsFuncWithOffset(offset int, diskNumber string) func() bool {
	return func() bool {
		return e.PartitionWithDataDirExistsWithOffset(offset+1, diskNumber)
	}
}

func (e *WindowsEnvironment) PartitionWithDataDirExistsFunc(diskNumber string) func() bool {
	return e.PartitionWithDataDirExistsFuncWithOffset(1, diskNumber)
}

func (e *WindowsEnvironment) EnsureVolumeHasDataDir(diskNumber string) {
	EventuallyWithOffset(1, e.PartitionWithDataDirExistsFunc(diskNumber), 2*time.Minute).Should(
		BeTrue(),
		fmt.Sprintf(`Expected partition with access path %s to be present on disk %s`, dataDir, diskNumber),
	)
}

func (e *WindowsEnvironment) EnsureAgentServiceStopped() {
	stdout := e.RunPowershellCommandWithOffset(1, "Get-Service -Name bosh-agent | Format-List -Property Status")

	running, err := regexp.MatchString("Running", stdout)
	ExpectWithOffset(1, err).NotTo(HaveOccurred())
	if running {
		e.RunPowershellCommandWithOffset(1, "c:\\bosh\\service_wrapper.exe stop")
	}
}

func (e *WindowsEnvironment) EnsureDataDirDoesntExist() {
	testPathOutput := e.RunPowershellCommandWithOffset(1, "Test-Path -Path %s", dataDir)

	exists := strings.TrimSpace(testPathOutput) == "True"
	if exists {
		e.RunPowershellCommandWithOffset(1, "Remove-Item %s -Force -Recurse", dataDir)
	}
}

func (e *WindowsEnvironment) AgentProcessRunningFunc() func() bool {
	return func() bool {
		exitCode, err := e.Client.Run(
			winrm.Powershell("Get-Process -ProcessName bosh-agent"),
			ioutil.Discard, ioutil.Discard,
		)
		return exitCode == 0 && err == nil
	}
}

func (e *WindowsEnvironment) RunPowershellCommandWithOffset(offset int, cmd string, cmdFmtArgs ...interface{}) string {
	outString, errString, exitCode, err := e.RunPowershellCommandWithOffsetAndResponses(offset+1, cmd, cmdFmtArgs...)

	ExpectWithOffset(offset+1, err).NotTo(
		HaveOccurred(),
		fmt.Sprintf(`Command "%s" failed with stdout: %s; stderr: %s`, cmd, outString, errString),
	)
	ExpectWithOffset(offset+1, exitCode).To(
		BeZero(),
		fmt.Sprintf(
			`Command "%s" failed with exit code: %d; stdout: %s; stderr: %s`, cmd, exitCode, outString, errString,
		),
	)

	return outString
}

func (e *WindowsEnvironment) RunPowershellCommandWithOffsetAndResponses(
	offset int,
	cmd string,
	cmdFmtArgs ...interface{},
) (string, string, int, error) {

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}

	exitCode, err := e.Client.Run(winrm.Powershell(fmt.Sprintf(cmd, cmdFmtArgs...)), stdout, stderr)

	outString := stdout.String()
	errString := stderr.String()

	return outString, errString, exitCode, err
}

func (e *WindowsEnvironment) RunPowershellCommand(cmd string, cmdFmtArgs ...interface{}) string {
	return e.RunPowershellCommandWithOffset(1, cmd, cmdFmtArgs...)
}

func (e *WindowsEnvironment) AssertDataACLed() {
	testFile := filepath.Join(dataDir + "testfile")

	e.RunPowershellCommandWithOffset(1, "echo 'content' >> %s", testFile)
	checkACLsOutput := e.RunPowershellCommandWithOffset(1, "Check-Acls %s", dataDir)
	aclErrsCount := strings.Count(checkACLsOutput, "Error")
	ExpectWithOffset(1, aclErrsCount == 0).To(
		BeTrue(),
		fmt.Sprintf("Expected data directory to have correct ACLs. Counted %d errors.", aclErrsCount),
	)
}

func (e *WindowsEnvironment) EnsureDiskClearedWithOffset(offset int, diskNumber string) {
	partitionCountOutput := e.RunPowershellCommandWithOffset(
		offset+1,
		"Get-Disk -Number %s | Select -ExpandProperty NumberOfPartitions",
		diskNumber,
	)

	partitionCount, err := strconv.Atoi(strings.TrimSpace(partitionCountOutput))
	ExpectWithOffset(offset+1, err).NotTo(HaveOccurred())
	if partitionCount > 0 {
		e.RunPowershellCommandWithOffset(offset+1, "Clear-Disk -Number %s -Confirm:$false -RemoveData", diskNumber)
	}
}

func (e *WindowsEnvironment) EnsureDiskCleared(diskNumber string) {
	e.EnsureDiskClearedWithOffset(1, diskNumber)
}
<commit_msg>Print to stdout, no unused variables<commit_after>package windows_test

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/masterzen/winrm"

	. "github.com/onsi/gomega"
)

const dataDir = `C:\var\vcap\data\`

type WindowsEnvironment struct {
	Client *winrm.Client
}

func (e *WindowsEnvironment) ShrinkRootPartition() {
	retryableFailure := "Resize-Partition : Size Not Supported"
	retryableError := "net/http: timeout awaiting response headers"

	cmd := "Get-Partition -DriveLetter C | Resize-Partition -Size $(Get-PartitionSupportedSize -DriveLetter C).SizeMin"

	for i := 0; i < 5; i++ {
		stdout, stderr, exitCode, err := e.RunPowershellCommandWithOffsetAndResponses(
			1,
			cmd,
		)

		if err != nil {
			if strings.Contains(err.Error(), retryableError) {
				fmt.Printf("WinRM timed out on attempt %d of 5, waiting 5 seconds to retry command: %s\n", i+1, cmd)
				time.Sleep(5 * time.Second)
				continue
			} else {
				ExpectWithOffset(1, err).NotTo(
					HaveOccurred(),
					fmt.Sprintf(`Command "%s" failed with stdout: %s; stderr: %s`, cmd, stdout, stderr),
				)
			}
		}

		if exitCode != 0 {
			if strings.Contains(stderr, retryableFailure) {
				fmt.Printf("Failed to shrink disk on attempt %d of 5, waiting 5 seconds to retry\n", i+1)
				time.Sleep(5 * time.Second)
				continue
			} else {
				ExpectWithOffset(1, exitCode).To(
					BeZero(),
					fmt.Sprintf(
						`Command "%s" failed with exit code: %d; stdout: %s; stderr: %s`, cmd, exitCode, stdout, stderr,
					),
				)
			}
		}

		break
	}
}

func (e *WindowsEnvironment) EnsureRootPartitionAtMaxSize() {
	freeSpaceOutput := e.RunPowershellCommandWithOffset(
		1,
		"Get-Disk $(Get-Partition -DriveLetter C | Select -ExpandProperty DiskNumber) | Select -ExpandProperty LargestFreeExtent",
	)

	freeSpace, err := strconv.Atoi(strings.TrimSpace(freeSpaceOutput))
	ExpectWithOffset(1, err).NotTo(HaveOccurred())

	if freeSpace > 0 {
		e.RunPowershellCommandWithOffset(
			1,
			"Resize-Partition -DriveLetter C -Size $(Get-PartitionSupportedSize -DriveLetter C).SizeMax",
		)
	}
}

func (e *WindowsEnvironment) GetDataDirPartitionNumber() string {
	return strings.TrimSpace(e.RunPowershellCommandWithOffset(
		1,
		fmt.Sprintf(`Get-Partition | Where AccessPaths -Contains "%s" | Select -ExpandProperty PartitionNumber`, dataDir),
	))
}

func (e *WindowsEnvironment) PartitionWithDataDirExists(diskNumber string) bool {
	return e.PartitionWithDataDirExistsWithOffset(1, diskNumber)
}

func (e *WindowsEnvironment) PartitionWithDataDirExistsWithOffset(offset int, diskNumber string) bool {
	stdout := e.RunPowershellCommandWithOffset(
		offset+1,
		fmt.Sprintf(
			`Get-Partition | where AccessPaths -Contains "%s" | Select -ExpandProperty DiskNumber`,
			dataDir,
		),
	)

	return strings.TrimSpace(stdout) == diskNumber
}

func (e *WindowsEnvironment) PartitionWithDataDirExistsFuncWithOffset(offset int, diskNumber string) func() bool {
	return func() bool {
		return e.PartitionWithDataDirExistsWithOffset(offset+1, diskNumber)
	}
}

func (e *WindowsEnvironment) PartitionWithDataDirExistsFunc(diskNumber string) func() bool {
	return e.PartitionWithDataDirExistsFuncWithOffset(1, diskNumber)
}

func (e *WindowsEnvironment) EnsureVolumeHasDataDir(diskNumber string) {
	EventuallyWithOffset(1, e.PartitionWithDataDirExistsFunc(diskNumber), 2*time.Minute).Should(
		BeTrue(),
		fmt.Sprintf(`Expected partition with access path %s to be present on disk %s`, dataDir, diskNumber),
	)
}

func (e *WindowsEnvironment) EnsureAgentServiceStopped() {
	stdout := e.RunPowershellCommandWithOffset(1, "Get-Service -Name bosh-agent | Format-List -Property Status")

	running, err := regexp.MatchString("Running", stdout)
	ExpectWithOffset(1, err).NotTo(HaveOccurred())
	if running {
		e.RunPowershellCommandWithOffset(1, "c:\\bosh\\service_wrapper.exe stop")
	}
}

func (e *WindowsEnvironment) EnsureDataDirDoesntExist() {
	testPathOutput := e.RunPowershellCommandWithOffset(1, "Test-Path -Path %s", dataDir)

	exists := strings.TrimSpace(testPathOutput) == "True"
	if exists {
		e.RunPowershellCommandWithOffset(1, "Remove-Item %s -Force -Recurse", dataDir)
	}
}

func (e *WindowsEnvironment) AgentProcessRunningFunc() func() bool {
	return func() bool {
		exitCode, err := e.Client.Run(
			winrm.Powershell("Get-Process -ProcessName bosh-agent"),
			ioutil.Discard, ioutil.Discard,
		)
		return exitCode == 0 && err == nil
	}
}

func (e *WindowsEnvironment) RunPowershellCommandWithOffset(offset int, cmd string, cmdFmtArgs ...interface{}) string {
	outString, errString, exitCode, err := e.RunPowershellCommandWithOffsetAndResponses(offset+1, cmd, cmdFmtArgs...)

	ExpectWithOffset(offset+1, err).NotTo(
		HaveOccurred(),
		fmt.Sprintf(`Command "%s" failed with stdout: %s; stderr: %s`, cmd, outString, errString),
	)
	ExpectWithOffset(offset+1, exitCode).To(
		BeZero(),
		fmt.Sprintf(
			`Command "%s" failed with exit code: %d; stdout: %s; stderr: %s`, cmd, exitCode, outString, errString,
		),
	)

	return outString
}

func (e *WindowsEnvironment) RunPowershellCommandWithOffsetAndResponses(
	offset int,
	cmd string,
	cmdFmtArgs ...interface{},
) (string, string, int, error) {

	stdout := &bytes.Buffer{}
	stderr := &bytes.Buffer{}

	exitCode, err := e.Client.Run(winrm.Powershell(fmt.Sprintf(cmd, cmdFmtArgs...)), stdout, stderr)

	outString := stdout.String()
	errString := stderr.String()

	return outString, errString, exitCode, err
}

func (e *WindowsEnvironment) RunPowershellCommand(cmd string, cmdFmtArgs ...interface{}) string {
	return e.RunPowershellCommandWithOffset(1, cmd, cmdFmtArgs...)
}

func (e *WindowsEnvironment) AssertDataACLed() {
	testFile := filepath.Join(dataDir + "testfile")

	e.RunPowershellCommandWithOffset(1, "echo 'content' >> %s", testFile)
	checkACLsOutput := e.RunPowershellCommandWithOffset(1, "Check-Acls %s", dataDir)
	aclErrsCount := strings.Count(checkACLsOutput, "Error")
	ExpectWithOffset(1, aclErrsCount == 0).To(
		BeTrue(),
		fmt.Sprintf("Expected data directory to have correct ACLs. Counted %d errors.", aclErrsCount),
	)
}

func (e *WindowsEnvironment) EnsureDiskClearedWithOffset(offset int, diskNumber string) {
	partitionCountOutput := e.RunPowershellCommandWithOffset(
		offset+1,
		"Get-Disk -Number %s | Select -ExpandProperty NumberOfPartitions",
		diskNumber,
	)

	partitionCount, err := strconv.Atoi(strings.TrimSpace(partitionCountOutput))
	ExpectWithOffset(offset+1, err).NotTo(HaveOccurred())
	if partitionCount > 0 {
		e.RunPowershellCommandWithOffset(offset+1, "Clear-Disk -Number %s -Confirm:$false -RemoveData", diskNumber)
	}
}

func (e *WindowsEnvironment) EnsureDiskCleared(diskNumber string) {
	e.EnsureDiskClearedWithOffset(1, diskNumber)
}
<|endoftext|>