<commit_before>package fgae

import(
	"fmt"
	"time"

	"github.com/skypies/geo/sfo"

	fdb "github.com/skypies/flightdb"
	"github.com/skypies/flightdb/ref"
)

// {{{ currentAccumulationTrack

func currentAccumulationTrack(f *fdb.Flight) *fdb.Track {
	if !f.HasTrack("ADSB") && !f.HasTrack("MLAT") { return nil }
	if !f.HasTrack("ADSB") { return f.Tracks["MLAT"] }
	if !f.HasTrack("MLAT") { return f.Tracks["ADSB"] }

	mlat,adsb := f.Tracks["MLAT"],f.Tracks["ADSB"]

	if len(*mlat) == 0 { return adsb }
	if len(*adsb) == 0 { return mlat }

	// Both tracks exist and are not empty ! Return most recent
	if (*mlat).End().After( (*adsb).End() ) {
		return mlat
	} else {
		return adsb
	}
}

// }}}
// {{{ AddTrackFragment

func (db *FlightDB)AddTrackFragment(frag *fdb.TrackFragment, airframes *ref.AirframeCache, schedules *ref.ScheduleCache) error {
	db.Debugf("* adding frag %d\n", len(frag.Track))
	f,err := db.LookupMostRecent(db.NewQuery().ByIcaoId(frag.IcaoId))
	if err != nil { return err }

	prefix := fmt.Sprintf("[%s/%s]%s %s", frag.IcaoId, frag.Callsign, frag.DataSystem, time.Now())

	// If the fragment is strictly a suffix, this will hold the preceding point
	var prevTP *fdb.Trackpoint

	if f == nil {
		f = fdb.NewFlightFromTrackFragment(frag)
		f.DebugLog += "-- AddFrag "+prefix+": new IcaoID\n"
		db.Debugf("* %s brand new IcaoID: %s", prefix, f)

	} else {
		db.Debugf("* %s found %s", prefix, f)

		trackKey := frag.TrackName() // ADSB, or MLAT; this is the track we will accumulate into

		// This is the most recent track we've accumulated into (could be ADSB, or MLAT); nil if none.
		// Note that this doesn't have to be the same as trackKey; we might be adding ADSB, but already
		// have some MLAT for the flight.
		accTrack := currentAccumulationTrack(f)

		if accTrack == nil {
			f.DebugLog += "-- AddFrag "+prefix+": first frag on pre-existing flight\n"
			db.Debugf("* %s no pre-existing track; adding right in", prefix)
			f.Tracks[trackKey] = &frag.Track

		} else if plausible,debug := accTrack.PlausibleContribution(&frag.Track); plausible==true {
			f.DebugLog += fmt.Sprintf("-- AddFrag %s: extending (adding %d to %d points)\n",
				prefix, len(frag.Track), len(*accTrack))
			db.Debugf("* %s extending track ... debug:\n%s", prefix, debug)

			// For MLAT data, callsigns can take a while to show up in the stream
			if f.Identity.Callsign == "" && frag.Callsign != "" {
				f.DebugLog += fmt.Sprintf(" - prev callsign was nil; adding it in now\n")
				f.Identity.Callsign = frag.Callsign
			}

			if !f.HasTrack(trackKey) {
				// If the accTrack was a different type (MLAT vs. ADSB), then we'll need to init
				f.Tracks[trackKey] = &fdb.Track{}

			} else {
				// Determine whether this frag is strictly a suffix to existing track data; this is the
				// common case. If so, keep a pointer to the trackpoint that precedes the frag
				n := len(*f.Tracks[trackKey])
				if n>0 && (*f.Tracks[trackKey])[n-1].TimestampUTC.Before(frag.Track[0].TimestampUTC) {
					db.Debugf("** new frag is strictly a suffix; prev = %d", n-1)
					prevTP = &((*f.Tracks[trackKey])[n-1])
				}
			}

			db.Debugf("* %s adding %d points to %d\n", prefix, len(frag.Track), len(*f.Tracks[trackKey]))

			db.Debugf("** pre : %s", f.Tracks[trackKey])
			f.Tracks[trackKey].Merge(&frag.Track)
			db.Debugf("** post: %s", f.Tracks[trackKey])

		}	else {
			f = fdb.NewFlightFromTrackFragment(frag)
			f.DebugLog += "-- AddFrag "+prefix+": was not plausible, so new flight\n"
			db.Debugf("* %s not a plausible addition; starting afresh ... debug\n%s", prefix, debug)
			f.DebugLog += debug+"\n"
		}
	}

	// Consult the airframe cache, and perhaps add some metadata, if not already present
	if f.Airframe.Registration == "" {
		if airframes != nil {
			if af := airframes.Get(f.IcaoId); af != nil {
				f.DebugLog += "-- AddFrag "+prefix+": found airframe\n"
				f.OverlayAirframe(*af)
			}
		}
	}

	// There could be a big gap between the previous track and this frag.
	// If that's the case, grab the preceding trackpoint and prefix this frag with it; then
	// the waypoint detection code (which builds lines between points) will look at the gap
	// between the frags, and maybe find extra waypoints.
	if prevTP != nil {
		// shift(x,a) : a = append([]T{x}, a...)
		frag.Track = append([]fdb.Trackpoint{*prevTP}, frag.Track...)
	}

	// Incrementally identify waypoints, frag by frag
	for wp,t := range frag.Track.MatchWaypoints(sfo.KFixes) {
		f.DebugLog += "-- AddFrag "+prefix+": found waypoint "+wp+"\n"
		f.SetWaypoint(wp,t)
	}

	return db.PersistFlight(f)
}

// }}}

// {{{ -------------------------={ E N D }=----------------------------------

// Local variables:
// folded-file: t
// end:

// }}}
<commit_msg>Add some perf tracing<commit_after>package fgae

import(
	"fmt"
	"time"

	"github.com/skypies/geo/sfo"

	fdb "github.com/skypies/flightdb"
	"github.com/skypies/flightdb/ref"
)

// {{{ currentAccumulationTrack

func currentAccumulationTrack(f *fdb.Flight) *fdb.Track {
	if !f.HasTrack("ADSB") && !f.HasTrack("MLAT") { return nil }
	if !f.HasTrack("ADSB") { return f.Tracks["MLAT"] }
	if !f.HasTrack("MLAT") { return f.Tracks["ADSB"] }

	mlat,adsb := f.Tracks["MLAT"],f.Tracks["ADSB"]

	if len(*mlat) == 0 { return adsb }
	if len(*adsb) == 0 { return mlat }

	// Both tracks exist and are not empty ! Return most recent
	if (*mlat).End().After( (*adsb).End() ) {
		return mlat
	} else {
		return adsb
	}
}

// }}}
// {{{ AddTrackFragment

func (db *FlightDB)AddTrackFragment(frag *fdb.TrackFragment, airframes *ref.AirframeCache, schedules *ref.ScheduleCache, perf map[string]time.Time) error {
	perf["01_start"] = time.Now()
	db.Debugf("* adding frag %d\n", len(frag.Track))

	f,err := db.LookupMostRecent(db.NewQuery().ByIcaoId(frag.IcaoId))
	if err != nil { return err }
	perf["02_mostrecent"] = time.Now()

	prefix := fmt.Sprintf("[%s/%s]%s %s", frag.IcaoId, frag.Callsign, frag.DataSystem, time.Now())

	// If the fragment is strictly a suffix, this will hold the preceding point
	var prevTP *fdb.Trackpoint

	if f == nil {
		f = fdb.NewFlightFromTrackFragment(frag)
		f.DebugLog += "-- AddFrag "+prefix+": new IcaoID\n"
		db.Debugf("* %s brand new IcaoID: %s", prefix, f)

	} else {
		db.Debugf("* %s found %s", prefix, f)

		trackKey := frag.TrackName() // ADSB, or MLAT; this is the track we will accumulate into

		// This is the most recent track we've accumulated into (could be ADSB, or MLAT); nil if none.
		// Note that this doesn't have to be the same as trackKey; we might be adding ADSB, but already
		// have some MLAT for the flight.
		accTrack := currentAccumulationTrack(f)

		if accTrack == nil {
			f.DebugLog += "-- AddFrag "+prefix+": first frag on pre-existing flight\n"
			db.Debugf("* %s no pre-existing track; adding right in", prefix)
			f.Tracks[trackKey] = &frag.Track

		} else if plausible,debug := accTrack.PlausibleContribution(&frag.Track); plausible==true {
			perf["03_plausible"] = time.Now()

			f.DebugLog += fmt.Sprintf("-- AddFrag %s: extending (adding %d to %d points)\n",
				prefix, len(frag.Track), len(*accTrack))
			db.Debugf("* %s extending track ... debug:\n%s", prefix, debug)

			// For MLAT data, callsigns can take a while to show up in the stream
			if f.Identity.Callsign == "" && frag.Callsign != "" {
				f.DebugLog += fmt.Sprintf(" - prev callsign was nil; adding it in now\n")
				f.Identity.Callsign = frag.Callsign
			}

			if !f.HasTrack(trackKey) {
				// If the accTrack was a different type (MLAT vs. ADSB), then we'll need to init
				f.Tracks[trackKey] = &fdb.Track{}

			} else {
				// Determine whether this frag is strictly a suffix to existing track data; this is the
				// common case. If so, keep a pointer to the trackpoint that precedes the frag
				n := len(*f.Tracks[trackKey])
				if n>0 && (*f.Tracks[trackKey])[n-1].TimestampUTC.Before(frag.Track[0].TimestampUTC) {
					db.Debugf("** new frag is strictly a suffix; prev = %d", n-1)
					prevTP = &((*f.Tracks[trackKey])[n-1])
				}
			}

			db.Debugf("* %s adding %d points to %d\n", prefix, len(frag.Track), len(*f.Tracks[trackKey]))

			db.Debugf("** pre : %s", f.Tracks[trackKey])
			f.Tracks[trackKey].Merge(&frag.Track)
			db.Debugf("** post: %s", f.Tracks[trackKey])

		}	else {
			perf["03_notplausible"] = time.Now()
			f = fdb.NewFlightFromTrackFragment(frag)
			f.DebugLog += "-- AddFrag "+prefix+": was not plausible, so new flight\n"
			db.Debugf("* %s not a plausible addition; starting afresh ... debug\n%s", prefix, debug)
			f.DebugLog += debug+"\n"
		}
	}

	// Consult the airframe cache, and perhaps add some metadata, if not already present
	if f.Airframe.Registration == "" {
		if airframes != nil {
			if af := airframes.Get(f.IcaoId); af != nil {
				f.DebugLog += "-- AddFrag "+prefix+": found airframe\n"
				f.OverlayAirframe(*af)
			}
		}
	}

	// There could be a big gap between the previous track and this frag.
	// If that's the case, grab the preceding trackpoint and prefix this frag with it; then
	// the waypoint detection code (which builds lines between points) will look at the gap
	// between the frags, and maybe find extra waypoints.
	if prevTP != nil {
		// shift(x,a) : a = append([]T{x}, a...)
		frag.Track = append([]fdb.Trackpoint{*prevTP}, frag.Track...)
	}

	perf["04_trackbuild"] = time.Now()

	// Incrementally identify waypoints, frag by frag
	for wp,t := range frag.Track.MatchWaypoints(sfo.KFixes) {
		f.DebugLog += "-- AddFrag "+prefix+": found waypoint "+wp+"\n"
		f.SetWaypoint(wp,t)
	}

	perf["05_waypoints"] = time.Now()
	err = db.PersistFlight(f)
	perf["06_persist"] = time.Now()

	return err
}

// }}}

// {{{ -------------------------={ E N D }=----------------------------------

// Local variables:
// folded-file: t
// end:

// }}}
<|endoftext|>


<commit_before>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package state

import (
	"context"
	"net/http"
	"net/url"

	"github.com/lxc/lxd/lxd/bgp"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/endpoints"
	"github.com/lxc/lxd/lxd/events"
	"github.com/lxc/lxd/lxd/firewall"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/maas"
	"github.com/lxc/lxd/lxd/sys"
	"github.com/lxc/lxd/shared"
)

// State is a gateway to the two main stateful components of LXD, the database
// and the operating system. It's typically used by model entities such as
// containers, volumes, etc. in order to perform changes.
type State struct {
	// Context
	Context context.Context

	// Databases
	Node    *db.Node
	Cluster *db.Cluster

	// MAAS server
	MAAS *maas.Controller

	// BGP server
	BGP *bgp.Server

	// OS access
	OS    *sys.OS
	Proxy func(req *http.Request) (*url.URL, error)

	// LXD server
	Endpoints *endpoints.Endpoints

	// Event server
	DevlxdEvents *events.Server
	Events       *events.Server

	// Firewall instance
	Firewall firewall.Firewall

	// Server certificate
	ServerCert             func() *shared.CertInfo
	UpdateCertificateCache func()

	// Available instance types based on operational drivers.
	InstanceTypes map[instancetype.Type]struct{}
}
<commit_msg>lxd/state: Add DevMonitor to state<commit_after>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package state

import (
	"context"
	"net/http"
	"net/url"

	"github.com/lxc/lxd/lxd/bgp"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/endpoints"
	"github.com/lxc/lxd/lxd/events"
	"github.com/lxc/lxd/lxd/firewall"
	"github.com/lxc/lxd/lxd/fsmonitor"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/maas"
	"github.com/lxc/lxd/lxd/sys"
	"github.com/lxc/lxd/shared"
)

// State is a gateway to the two main stateful components of LXD, the database
// and the operating system. It's typically used by model entities such as
// containers, volumes, etc. in order to perform changes.
type State struct {
	// Context
	Context context.Context

	// Databases
	Node    *db.Node
	Cluster *db.Cluster

	// MAAS server
	MAAS *maas.Controller

	// BGP server
	BGP *bgp.Server

	// OS access
	OS    *sys.OS
	Proxy func(req *http.Request) (*url.URL, error)

	// LXD server
	Endpoints *endpoints.Endpoints

	// Event server
	DevlxdEvents *events.Server
	Events       *events.Server

	// Firewall instance
	Firewall firewall.Firewall

	// Server certificate
	ServerCert             func() *shared.CertInfo
	UpdateCertificateCache func()

	// Available instance types based on operational drivers.
	InstanceTypes map[instancetype.Type]struct{}

	// Filesystem monitor
	DevMonitor fsmonitor.FSMonitor
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-client-go/index"
	"github.com/taskcluster/taskcluster-client-go/queue"
	"os"
	"regexp"
	"testing"
	"time"
)

// This is a silly test that looks for the latest mozilla-central buildbot linux64 l10n build
// and asserts that it must have a created time between a year ago and an hour in the future.
//
// Could easily break at a point in the future, at which point we can change to something else.
//
// Note, no credentials are needed, so this can be run even on travis-ci.org, for example.
func TestFindLatestBuildbotTask(t *testing.T) {
	Index := index.New("", "")
	Queue := queue.New("", "")
	itr, cs1 := Index.FindTask("buildbot.branches.mozilla-central.linux64.l10n")
	if cs1.Error != nil {
		t.Fatalf("%v\n", cs1.Error)
	}
	taskId := itr.TaskId
	td, cs2 := Queue.Task(taskId)
	if cs2.Error != nil {
		t.Fatalf("%v\n", cs2.Error)
	}
	created := time.Time(td.Created).Local()

	// calculate time an hour in the future to allow for clock drift
	now := time.Now().Local()
	inAnHour := now.Add(time.Hour * 1)
	aYearAgo := now.AddDate(-1, 0, 0)
	t.Log("")
	t.Log("  => Task " + taskId + " was created on " + created.Format("Mon, 2 Jan 2006 at 15:04:00 -0700"))
	t.Log("")
	if created.After(inAnHour) {
		t.Log("Current time: " + now.Format("Mon, 2 Jan 2006 at 15:04:00 -0700"))
		t.Error("Task " + taskId + " has a creation date that is over an hour in the future")
	}
	if created.Before(aYearAgo) {
		t.Log("Current time: " + now.Format("Mon, 2 Jan 2006 at 15:04:00 -0700"))
		t.Error("Task " + taskId + " has a creation date that is over a year old")
	}

}

// Tests whether it is possible to define a task against the production Queue.
func TestDefineTask(t *testing.T) {
	clientId := os.Getenv("TASKCLUSTER_CLIENT_ID")
	accessToken := os.Getenv("TASKCLUSTER_ACCESS_TOKEN")
	certificate := os.Getenv("TASKCLUSTER_CERTIFICATE")
	if clientId == "" || accessToken == "" {
		t.Skip("Skipping test TestDefineTask since TASKCLUSTER_CLIENT_ID and/or TASKCLUSTER_ACCESS_TOKEN env vars not set")
	}
	myQueue := queue.New(clientId, accessToken)
	myQueue.Certificate = certificate

	taskId := slugid.Nice()
	created := time.Now()
	deadline := created.AddDate(0, 0, 1)
	expires := deadline

	td := &queue.TaskDefinition{
		Created:  queue.Time(created),
		Deadline: queue.Time(deadline),
		Expires:  queue.Time(expires),
		Extra:    json.RawMessage(`{"index":{"rank":12345}}`),
		Metadata: struct {
			Description string `json:"description"`
			Name        string `json:"name"`
			Owner       string `json:"owner"`
			Source      string `json:"source"`
		}{
			Description: "Stuff",
			Name:        "[TC] Pete",
			Owner:       "pmoore@mozilla.com",
			Source:      "http://everywhere.com/",
		},
		Payload:       json.RawMessage(`{"features":{"relengApiProxy":true}}`),
		ProvisionerId: "win-provisioner",
		Retries:       5,
		Routes: []string{
			"tc-treeherder.mozilla-inbound.bcf29c305519d6e120b2e4d3b8aa33baaf5f0163",
			"tc-treeherder-stage.mozilla-inbound.bcf29c305519d6e120b2e4d3b8aa33baaf5f0163",
		},
		SchedulerId: "go-test-test-scheduler",
		Scopes: []string{
			"docker-worker:image:taskcluster/builder:0.5.6",
			"queue:define-task:aws-provisioner-v1/build-c4-2xlarge",
		},
		Tags:        json.RawMessage(`{"createdForUser":"cbook@mozilla.com"}`),
		Priority:    json.RawMessage(`"high"`),
		TaskGroupId: "dtwuF2n9S-i83G37V9eBuQ",
		WorkerType:  "win2008-worker",
	}

	tsr, cs := myQueue.DefineTask(taskId, td)

	//////////////////////////////////
	// And now validate results.... //
	//////////////////////////////////

	if cs.Error != nil {
		b := bytes.Buffer{}
		cs.HttpRequest.Header.Write(&b)
		headers := regexp.MustCompile(`(mac|nonce)="[^"]*"`).ReplaceAllString(b.String(), `$1="***********"`)
		t.Logf("\n\nRequest sent:\n\nURL: %s\nMethod: %s\nHeaders:\n%v\nBody: %s", cs.HttpRequest.URL, cs.HttpRequest.Method, headers, cs.HttpRequestBody)
		t.Fatalf("\n\nResponse received:\n\n%s", cs.Error)
	}

	t.Logf("Task https://queue.taskcluster.net/v1/task/%v created successfully", taskId)

	if provisionerId := cs.HttpRequestObject.(*queue.TaskDefinition).ProvisionerId; provisionerId != "win-provisioner" {
		t.Errorf("provisionerId 'win-provisioner' expected but got %s", provisionerId)
	}
	if schedulerId := tsr.Status.SchedulerId; schedulerId != "go-test-test-scheduler" {
		t.Errorf("schedulerId 'go-test-test-scheduler' expected but got %s", schedulerId)
	}
	if retriesLeft := tsr.Status.RetriesLeft; retriesLeft != 5 {
		t.Errorf("Expected 'retriesLeft' to be 5, but got %v", retriesLeft)
	}
	if state := tsr.Status.State; state != "unscheduled" {
		t.Errorf("Expected 'state' to be 'unscheduled', but got %s", state)
	}
	submittedPayload := cs.HttpRequestBody

	// only the contents is relevant below - the formatting and order of properties does not matter
	// since a json comparison is done, not a string comparison...
	expectedJson := []byte(`
	{
	  "created":  "` + created.UTC().Format("2006-01-02T15:04:05.000Z") + `",
	  "deadline": "` + deadline.UTC().Format("2006-01-02T15:04:05.000Z") + `",
	  "expires":  "` + expires.UTC().Format("2006-01-02T15:04:05.000Z") + `",

	  "taskGroupId": "dtwuF2n9S-i83G37V9eBuQ",
	  "workerType":  "win2008-worker",
	  "schedulerId": "go-test-test-scheduler",

	  "payload": {
	    "features": {
	      "relengApiProxy":true
	    }
	  },

	  "priority":      "high",
	  "provisionerId": "win-provisioner",
	  "retries":       5,

	  "routes": [
	    "tc-treeherder.mozilla-inbound.bcf29c305519d6e120b2e4d3b8aa33baaf5f0163",
	    "tc-treeherder-stage.mozilla-inbound.bcf29c305519d6e120b2e4d3b8aa33baaf5f0163"
	  ],

	  "scopes": [
	    "docker-worker:image:taskcluster/builder:0.5.6",
	    "queue:define-task:aws-provisioner-v1/build-c4-2xlarge"
	  ],

	  "tags": {
	    "createdForUser": "cbook@mozilla.com"
	  },

	  "extra": {
	    "index": {
	      "rank": 12345
	    }
	  },

	  "metadata": {
	    "description": "Stuff",
	    "name":        "[TC] Pete",
	    "owner":       "pmoore@mozilla.com",
	    "source":      "http://everywhere.com/"
	  }
	}
	`)

	jsonCorrect, formattedExpected, formattedActual, err := jsonEqual(expectedJson, []byte(submittedPayload))
	if err != nil {
		t.Fatalf("Exception thrown formatting json data!\n%s", err)
	}

	if !jsonCorrect {
		t.Log("Anticipated json not generated. Expected:")
		t.Logf("%s", formattedExpected)
		t.Log("Actual:")
		t.Errorf("%s", formattedActual)
	}
}

// Checks whether two json []byte are equivalent (equal) by formatting/ordering
// both of them consistently, and then comparing if formatted versions are
// identical. Returns true/false together with formatted json, and any error.
func jsonEqual(a []byte, b []byte) (bool, []byte, []byte, error) {
	a_, err := formatJson(a)
	if err != nil {
		return false, nil, nil, err
	}
	b_, err := formatJson(b)
	if err != nil {
		return false, a_, nil, err
	}
	return string(a_) == string(b_), a_, b_, nil
}

// Takes json []byte input, unmarshals and then marshals, in order to get a
// canonical representation of json (i.e. formatted with objects ordered).
// Ugly and perhaps inefficient, but effective! :p
func formatJson(a []byte) ([]byte, error) {
	tmpObj := new(interface{})
	err := json.Unmarshal(a, &tmpObj)
	if err != nil {
		return a, err
	}
	return json.MarshalIndent(&tmpObj, "", "  ")
}
<commit_msg>Better logging in tests<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"github.com/taskcluster/slugid-go/slugid"
	"github.com/taskcluster/taskcluster-client-go/index"
	"github.com/taskcluster/taskcluster-client-go/queue"
	"os"
	"regexp"
	"testing"
	"time"
)

// This is a silly test that looks for the latest mozilla-central buildbot linux64 l10n build
// and asserts that it must have a created time between a year ago and an hour in the future.
//
// Could easily break at a point in the future, at which point we can change to something else.
//
// Note, no credentials are needed, so this can be run even on travis-ci.org, for example.
func TestFindLatestBuildbotTask(t *testing.T) {
	Index := index.New("", "")
	Queue := queue.New("", "")
	itr, cs1 := Index.FindTask("buildbot.branches.mozilla-central.linux64.l10n")
	if cs1.Error != nil {
		t.Fatalf("%v\n", cs1.Error)
	}
	taskId := itr.TaskId
	td, cs2 := Queue.Task(taskId)
	if cs2.Error != nil {
		t.Fatalf("%v\n", cs2.Error)
	}
	created := time.Time(td.Created).Local()

	// calculate time an hour in the future to allow for clock drift
	now := time.Now().Local()
	inAnHour := now.Add(time.Hour * 1)
	aYearAgo := now.AddDate(-1, 0, 0)
	t.Log("")
	t.Log("  => Task " + taskId + " was created on " + created.Format("Mon, 2 Jan 2006 at 15:04:00 -0700"))
	t.Log("")
	if created.After(inAnHour) {
		t.Log("Current time: " + now.Format("Mon, 2 Jan 2006 at 15:04:00 -0700"))
		t.Error("Task " + taskId + " has a creation date that is over an hour in the future")
	}
	if created.Before(aYearAgo) {
		t.Log("Current time: " + now.Format("Mon, 2 Jan 2006 at 15:04:00 -0700"))
		t.Error("Task " + taskId + " has a creation date that is over a year old")
	}

}

// Tests whether it is possible to define a task against the production Queue.
func TestDefineTask(t *testing.T) {
	clientId := os.Getenv("TASKCLUSTER_CLIENT_ID")
	accessToken := os.Getenv("TASKCLUSTER_ACCESS_TOKEN")
	certificate := os.Getenv("TASKCLUSTER_CERTIFICATE")
	if clientId == "" || accessToken == "" {
		t.Skip("Skipping test TestDefineTask since TASKCLUSTER_CLIENT_ID and/or TASKCLUSTER_ACCESS_TOKEN env vars not set")
	}
	myQueue := queue.New(clientId, accessToken)
	myQueue.Certificate = certificate

	taskId := slugid.Nice()
	created := time.Now()
	deadline := created.AddDate(0, 0, 1)
	expires := deadline

	td := &queue.TaskDefinition{
		Created:  queue.Time(created),
		Deadline: queue.Time(deadline),
		Expires:  queue.Time(expires),
		Extra:    json.RawMessage(`{"index":{"rank":12345}}`),
		Metadata: struct {
			Description string `json:"description"`
			Name        string `json:"name"`
			Owner       string `json:"owner"`
			Source      string `json:"source"`
		}{
			Description: "Stuff",
			Name:        "[TC] Pete",
			Owner:       "pmoore@mozilla.com",
			Source:      "http://everywhere.com/",
		},
		Payload:       json.RawMessage(`{"features":{"relengApiProxy":true}}`),
		ProvisionerId: "win-provisioner",
		Retries:       5,
		Routes: []string{
			"tc-treeherder.mozilla-inbound.bcf29c305519d6e120b2e4d3b8aa33baaf5f0163",
			"tc-treeherder-stage.mozilla-inbound.bcf29c305519d6e120b2e4d3b8aa33baaf5f0163",
		},
		SchedulerId: "go-test-test-scheduler",
		Scopes: []string{
			"test-worker:image:toastposter/pumpkin:0.5.6",
		},
		Tags:        json.RawMessage(`{"createdForUser":"cbook@mozilla.com"}`),
		Priority:    json.RawMessage(`"high"`),
		TaskGroupId: "dtwuF2n9S-i83G37V9eBuQ",
		WorkerType:  "win2008-worker",
	}

	tsr, cs := myQueue.DefineTask(taskId, td)

	//////////////////////////////////
	// And now validate results.... //
	//////////////////////////////////

	if cs.Error != nil {
		b := bytes.Buffer{}
		cs.HttpRequest.Header.Write(&b)
		headers := regexp.MustCompile(`(mac|nonce)="[^"]*"`).ReplaceAllString(b.String(), `$1="***********"`)
		t.Logf("\n\nRequest sent:\n\nURL: %s\nMethod: %s\nHeaders:\n%v\nBody: %s", cs.HttpRequest.URL, cs.HttpRequest.Method, headers, cs.HttpRequestBody)
		t.Fatalf("\n\nResponse received:\n\n%s", cs.Error)
	}

	t.Logf("Task https://queue.taskcluster.net/v1/task/%v created successfully", taskId)

	if provisionerId := cs.HttpRequestObject.(*queue.TaskDefinition).ProvisionerId; provisionerId != "win-provisioner" {
		t.Errorf("provisionerId 'win-provisioner' expected but got %s", provisionerId)
	}
	if schedulerId := tsr.Status.SchedulerId; schedulerId != "go-test-test-scheduler" {
		t.Errorf("schedulerId 'go-test-test-scheduler' expected but got %s", schedulerId)
	}
	if retriesLeft := tsr.Status.RetriesLeft; retriesLeft != 5 {
		t.Errorf("Expected 'retriesLeft' to be 5, but got %v", retriesLeft)
	}
	if state := tsr.Status.State; state != "unscheduled" {
		t.Errorf("Expected 'state' to be 'unscheduled', but got %s", state)
	}
	submittedPayload := cs.HttpRequestBody

	// only the contents is relevant below - the formatting and order of properties does not matter
	// since a json comparison is done, not a string comparison...
	expectedJson := []byte(`
	{
	  "created":  "` + created.UTC().Format("2006-01-02T15:04:05.000Z") + `",
	  "deadline": "` + deadline.UTC().Format("2006-01-02T15:04:05.000Z") + `",
	  "expires":  "` + expires.UTC().Format("2006-01-02T15:04:05.000Z") + `",

	  "taskGroupId": "dtwuF2n9S-i83G37V9eBuQ",
	  "workerType":  "win2008-worker",
	  "schedulerId": "go-test-test-scheduler",

	  "payload": {
	    "features": {
	      "relengApiProxy":true
	    }
	  },

	  "priority":      "high",
	  "provisionerId": "win-provisioner",
	  "retries":       5,

	  "routes": [
	    "tc-treeherder.mozilla-inbound.bcf29c305519d6e120b2e4d3b8aa33baaf5f0163",
	    "tc-treeherder-stage.mozilla-inbound.bcf29c305519d6e120b2e4d3b8aa33baaf5f0163"
	  ],

	  "scopes": [
	    "test-worker:image:toastposter/pumpkin:0.5.6"
	  ],

	  "tags": {
	    "createdForUser": "cbook@mozilla.com"
	  },

	  "extra": {
	    "index": {
	      "rank": 12345
	    }
	  },

	  "metadata": {
	    "description": "Stuff",
	    "name":        "[TC] Pete",
	    "owner":       "pmoore@mozilla.com",
	    "source":      "http://everywhere.com/"
	  }
	}
	`)

	jsonCorrect, formattedExpected, formattedActual, err := jsonEqual(expectedJson, []byte(submittedPayload))
	if err != nil {
		t.Fatalf("Exception thrown formatting json data!\n%s\n\nStruggled to format either:\n%s\n\nor:\n\n%s", err, string(expectedJson), submittedPayload)
	}

	if !jsonCorrect {
		t.Log("Anticipated json not generated. Expected:")
		t.Logf("%s", formattedExpected)
		t.Log("Actual:")
		t.Errorf("%s", formattedActual)
	}
}

// Checks whether two json []byte are equivalent (equal) by formatting/ordering
// both of them consistently, and then comparing if formatted versions are
// identical. Returns true/false together with formatted json, and any error.
func jsonEqual(a []byte, b []byte) (bool, []byte, []byte, error) {
	a_, err := formatJson(a)
	if err != nil {
		return false, nil, nil, err
	}
	b_, err := formatJson(b)
	if err != nil {
		return false, a_, nil, err
	}
	return string(a_) == string(b_), a_, b_, nil
}

// Takes json []byte input, unmarshals and then marshals, in order to get a
// canonical representation of json (i.e. formatted with objects ordered).
// Ugly and perhaps inefficient, but effective! :p
func formatJson(a []byte) ([]byte, error) {
	tmpObj := new(interface{})
	err := json.Unmarshal(a, &tmpObj)
	if err != nil {
		return a, err
	}
	return json.MarshalIndent(&tmpObj, "", "  ")
}
<|endoftext|>


<commit_before>package lb

import (
	"fmt"
	"strings"
	"time"

	"golang.org/x/net/context"

	"github.com/go-kit/kit/endpoint"
)

// Callback function that indicates the current attempt count and the error encountered.
//  Should return whether the Retry function should continue trying, and a custom
//  error message if desired. The error message may be nil, but a true/false
//  is always expected. In all cases if the error message is supplied, the
//  current error will be replaced.
type Callback func(n int, err error) (cont bool, cbErr error)

// Retry wraps a service load balancer and returns an endpoint oriented load
// balancer for the specified service method.
// Requests to the endpoint will be automatically load balanced via the load
// balancer. Requests that return errors will be retried until they succeed,
// up to max times, or until the timeout is elapsed, whichever comes first.
func Retry(max int, timeout time.Duration, b Balancer) endpoint.Endpoint {
	return RetryWithCallback(max, timeout, b, func(c int, err error) (bool, error) { return true, nil })
}

// RetryWithCallback wraps a service load balancer and returns an endpoint oriented load
// balancer for the specified service method.
// Requests to the endpoint will be automatically load balanced via the load
// balancer. Requests that return errors will be retried until they succeed,
// up to max times, until the callback returns false, or until the timeout is elapsed,
// whichever comes first.
func RetryWithCallback(max int, timeout time.Duration, b Balancer, cb Callback) endpoint.Endpoint {
	if cb == nil {
		panic("nil Callback")
	}
	if b == nil {
		panic("nil Balancer")
	}
	return func(ctx context.Context, request interface{}) (response interface{}, err error) {
		var (
			newctx, cancel = context.WithTimeout(ctx, timeout)
			responses      = make(chan interface{}, 1)
			errs           = make(chan error, 1)
			a              = []string{}
		)
		defer cancel()
		for i := 1; i <= max; i++ {
			go func() {
				e, err := b.Endpoint()
				if err != nil {
					errs <- err
					return
				}
				response, err := e(newctx, request)
				if err != nil {
					errs <- err
					return
				}
				responses <- response
			}()

			select {
			case <-newctx.Done():
				return nil, newctx.Err()
			case response := <-responses:
				return response, nil
			case err := <-errs:
				cont, cbErr := cb(i, err)
				if !cont {
					if cbErr == nil {
						return nil, fmt.Errorf("retry attempts exceeded (%s)", strings.Join(a, "; "))
					}
					return nil, cbErr
				}
				currentErr := err.Error()
				if cbErr != nil {
					currentErr = cbErr.Error()
				}
				a = append(a, currentErr)
				continue
			}
		}
		return nil, fmt.Errorf("retry attempts exceeded (%s)", strings.Join(a, "; "))
	}
}
<commit_msg>Tidy formatting of callback comment.<commit_after>package lb

import (
	"fmt"
	"strings"
	"time"

	"golang.org/x/net/context"

	"github.com/go-kit/kit/endpoint"
)

// Callback is a function that indicates the current attempt count and the error
// encountered. Should return whether the Retry function should continue trying,
// and a custom error message if desired. The error message may be nil, but a
// true/false is always expected. In all cases if the error message is supplied,
// thecurrent error will be replaced.
type Callback func(n int, err error) (cont bool, cbErr error)

// Retry wraps a service load balancer and returns an endpoint oriented load
// balancer for the specified service method.
// Requests to the endpoint will be automatically load balanced via the load
// balancer. Requests that return errors will be retried until they succeed,
// up to max times, or until the timeout is elapsed, whichever comes first.
func Retry(max int, timeout time.Duration, b Balancer) endpoint.Endpoint {
	return RetryWithCallback(max, timeout, b, func(c int, err error) (bool, error) { return true, nil })
}

// RetryWithCallback wraps a service load balancer and returns an endpoint oriented load
// balancer for the specified service method.
// Requests to the endpoint will be automatically load balanced via the load
// balancer. Requests that return errors will be retried until they succeed,
// up to max times, until the callback returns false, or until the timeout is elapsed,
// whichever comes first.
func RetryWithCallback(max int, timeout time.Duration, b Balancer, cb Callback) endpoint.Endpoint {
	if cb == nil {
		panic("nil Callback")
	}
	if b == nil {
		panic("nil Balancer")
	}
	return func(ctx context.Context, request interface{}) (response interface{}, err error) {
		var (
			newctx, cancel = context.WithTimeout(ctx, timeout)
			responses      = make(chan interface{}, 1)
			errs           = make(chan error, 1)
			a              = []string{}
		)
		defer cancel()
		for i := 1; i <= max; i++ {
			go func() {
				e, err := b.Endpoint()
				if err != nil {
					errs <- err
					return
				}
				response, err := e(newctx, request)
				if err != nil {
					errs <- err
					return
				}
				responses <- response
			}()

			select {
			case <-newctx.Done():
				return nil, newctx.Err()
			case response := <-responses:
				return response, nil
			case err := <-errs:
				cont, cbErr := cb(i, err)
				if !cont {
					if cbErr == nil {
						return nil, fmt.Errorf("retry attempts exceeded (%s)", strings.Join(a, "; "))
					}
					return nil, cbErr
				}
				currentErr := err.Error()
				if cbErr != nil {
					currentErr = cbErr.Error()
				}
				a = append(a, currentErr)
				continue
			}
		}
		return nil, fmt.Errorf("retry attempts exceeded (%s)", strings.Join(a, "; "))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/koding/kite"
	"fmt"
	"os"
	"github.com/bitly/go-nsq"
)

func main() {
	k := kite.New("robber", "1.0.0")
	k.Config.Port = 6002
	k.Config.DisableAuthentication = true
	k.HandleFunc("robber", robber)

	fmt.Println("Robber staring kite")
	k.Run()

	fmt.Println("Robber configuring NSQ")
	wg := &sync.WaitGroup{}
	wg.Add(1)

	config := nsq.NewConfig()
	q, _ := nsq.NewConsumer("tick", "ch", config)
	q.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
		log.Printf("Got a message: %v", message)
		wg.Done()
		return nil
	}))

	err := q.ConnectToNSQD(os.Getenv("MESSAGING_SERVICE_HOST" + ":4150"))
	if err != nil {
		log.Panic("Could not connect")
	}
	fmt.Println("Robber starting NSQ")
	wg.Wait()

}

func hello(r *kite.Request) (interface{}, error) {

	fmt.Println("Robber got hello")

	// You can return anything as result, as long as it is JSON marshalable.
	return nil
}
<commit_msg>syntax<commit_after>package main

import (
	"github.com/koding/kite"
	"fmt"
	"os"
	"sync"
	"log"
	"github.com/bitly/go-nsq"
)

func main() {
	k := kite.New("robber", "1.0.0")
	k.Config.Port = 6002
	k.Config.DisableAuthentication = true
	k.HandleFunc("hello", hello)

	fmt.Println("Robber staring kite")
	k.Run()

	fmt.Println("Robber configuring NSQ")
	wg := &sync.WaitGroup{}
	wg.Add(1)

	config := nsq.NewConfig()
	q, _ := nsq.NewConsumer("tick", "ch", config)
	q.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
		log.Printf("Got a message: %v", message)
		wg.Done()
		return nil
	}))

	err := q.ConnectToNSQD(os.Getenv("MESSAGING_SERVICE_HOST" + ":4150"))
	if err != nil {
		log.Panic("Could not connect")
	}
	fmt.Println("Robber starting NSQ")
	wg.Wait()

}

func hello(r *kite.Request) (interface{}, error) {

	fmt.Println("Robber got hello")

	// You can return anything as result, as long as it is JSON marshalable.
	return nil
}
<|endoftext|>


<commit_before>package rev

import (
	"fmt"
	"net/http"
	"net/url"
	"regexp"
	"testing"
)

// Data-driven tests that check that a given routes-file line translates into
// the expected Route object.
var routeTestCases = map[string]*Route{
	"get / Application.Index": &Route{
		method:        "GET",
		path:          "/",
		action:        "Application.Index",
		pathPattern:   regexp.MustCompile("/$"),
		staticDir:     "",
		args:          []*arg{},
		actionPattern: regexp.MustCompile("Application\\.Index"),
	},

	"post /app/{id} Application.SaveApp": &Route{
		method:      "POST",
		path:        "/app/{id}",
		action:      "Application.SaveApp",
		pathPattern: regexp.MustCompile("/app/(?P<id>[^/]+)$"),
		staticDir:   "",
		args: []*arg{
			{
				name:       "id",
				constraint: regexp.MustCompile("[^/]+"),
			},
		},
		actionPattern: regexp.MustCompile("Application\\.SaveApp"),
	},

	"GET /public/ staticDir:www": &Route{
		method:        "GET",
		path:          "/public/",
		action:        "staticDir:www",
		pathPattern:   regexp.MustCompile("^/public/(.*)$"),
		staticDir:     "www",
		args:          []*arg{},
		actionPattern: nil,
	},

	"* /{controller}/{action} {controller}.{action}": &Route{
		method:      "*",
		path:        "/{controller}/{action}",
		action:      "{controller}.{action}",
		pathPattern: regexp.MustCompile("/(?P<controller>[^/]+)/(?P<action>[^/]+)$"),
		staticDir:   "",
		args: []*arg{
			{
				name:       "controller",
				constraint: regexp.MustCompile("[^/]+"),
			},
			{
				name:       "action",
				constraint: regexp.MustCompile("[^/]+"),
			},
		},
		actionPattern: regexp.MustCompile("(?P<controller>[^/]+)\\.(?P<action>[^/]+)"),
		actionArgs:    []string{"controller", "action"},
	},
}

// Run the test cases above.
func TestComputeRoute(t *testing.T) {
	for routeLine, expected := range routeTestCases {
		method, path, action, found := parseRouteLine(routeLine)
		if !found {
			t.Error("Failed to parse route line:", routeLine)
			continue
		}
		actual := NewRoute(method, path, action)
		eq(t, "Method", actual.method, expected.method)
		eq(t, "Path", actual.path, expected.path)
		eq(t, "Action", actual.action, expected.action)
		eq(t, "pathPattern", fmt.Sprint(actual.pathPattern), fmt.Sprint(expected.pathPattern))
		eq(t, "staticDir", actual.staticDir, expected.staticDir)
		eq(t, "len(args)", len(actual.args), len(expected.args))
		for i, arg := range actual.args {
			if len(expected.args) <= i {
				break
			}
			eq(t, "arg.name", arg.name, expected.args[i].name)
			eq(t, "arg.constraint", arg.constraint.String(), expected.args[i].constraint.String())
		}
		eq(t, "actionPattern", fmt.Sprint(actual.actionPattern), fmt.Sprint(expected.actionPattern))
		eq(t, "len(actionArgs)", len(actual.actionArgs), len(expected.actionArgs))
		eq(t, "actionArgs", fmt.Sprint(actual.actionArgs), fmt.Sprint(expected.actionArgs))
		if t.Failed() {
			t.Fatal("Failed on route:", routeLine)
		}
	}
}

// Router Tests

const TEST_ROUTES = `
# This is a comment
GET  /                       Application.Index
GET  /app/{id}               Application.Show
POST /app/{id}               Application.Save

GET  /public/                staticDir:www
*    /{controller}/{action} {controller}.{action}
`

var routeMatchTestCases = map[*http.Request]*RouteMatch{
	&http.Request{
		Method: "GET",
		URL:    &url.URL{Path: "/"},
	}: &RouteMatch{
		ControllerName: "Application",
		MethodName:     "Index",
		Params:         map[string]string{},
		StaticFilename: "",
	},

	&http.Request{
		Method: "GET",
		URL:    &url.URL{Path: "/app/123"},
	}: &RouteMatch{
		ControllerName: "Application",
		MethodName:     "ShowApp",
		Params:         map[string]string{"id": "123"},
		StaticFilename: "",
	},

	&http.Request{
		Method: "POST",
		URL:    &url.URL{Path: "/app/123"},
	}: &RouteMatch{
		ControllerName: "Application",
		MethodName:     "SaveApp",
		Params:         map[string]string{"id": "123"},
		StaticFilename: "",
	},

	&http.Request{
		Method: "GET",
		URL:    &url.URL{Path: "/public/style.css"},
	}: &RouteMatch{
		ControllerName: "",
		MethodName:     "",
		Params:         map[string]string{},
		StaticFilename: "www/style.css",
	},

	&http.Request{
		Method: "GET",
		URL:    &url.URL{Path: "/Implicit/Route"},
	}: &RouteMatch{
		ControllerName: "Implicit",
		MethodName:     "Route",
		Params:         map[string]string{"controller": "Implicit", "action": "Route"},
		StaticFilename: "",
	},
}

func TestRouteMatches(t *testing.T) {
	router := NewRouter(TEST_ROUTES)
	for req, expected := range routeMatchTestCases {
		actual := router.Route(req)
		if !eq(t, "Found route", actual != nil, expected != nil) {
			continue
		}
		eq(t, "ControllerName", actual.ControllerName, expected.ControllerName)
		eq(t, "MethodName", actual.ControllerName, expected.ControllerName)
		eq(t, "len(Params)", len(actual.Params), len(expected.Params))
		for key, actualValue := range actual.Params {
			eq(t, "Params", actualValue, expected.Params[key])
		}
		eq(t, "StaticFilename", actual.StaticFilename, expected.StaticFilename)
	}
}

// Reverse Routing

type ReverseRouteArgs struct {
	action string
	args   map[string]string
}

var reverseRoutingTestCases = map[*ReverseRouteArgs]*ActionDefinition{
	&ReverseRouteArgs{
		action: "Application.Index",
		args:   map[string]string{},
	}: &ActionDefinition{
		Url:    "/",
		Method: "GET",
		Star:   false,
		Action: "Application.Index",
	},

	&ReverseRouteArgs{
		action: "Application.Show",
		args:   map[string]string{"id": "123"},
	}: &ActionDefinition{
		Url:    "/app/123",
		Method: "GET",
		Star:   false,
		Action: "Application.Show",
	},

	&ReverseRouteArgs{
		action: "Implicit.Route",
		args:   map[string]string{},
	}: &ActionDefinition{
		Url:    "/Implicit/Route",
		Method: "GET",
		Star:   true,
		Action: "Implicit.Route",
	},

	&ReverseRouteArgs{
		action: "Application.Save",
		args:   map[string]string{"id": "123", "c": "http://continue"},
	}: &ActionDefinition{
		Url:    "/app/123?c=http%3A%2F%2Fcontinue",
		Method: "POST",
		Star:   false,
		Action: "Application.Save",
	},
}

func TestReverseRouting(t *testing.T) {
	router := NewRouter(TEST_ROUTES)
	for routeArgs, expected := range reverseRoutingTestCases {
		actual := router.Reverse(routeArgs.action, routeArgs.args)
		if !eq(t, "Found route", actual != nil, expected != nil) {
			continue
		}
		eq(t, "Url", actual.Url, expected.Url)
		eq(t, "Method", actual.Method, expected.Method)
		eq(t, "Star", actual.Star, expected.Star)
		eq(t, "Action", actual.Action, expected.Action)
	}
}

// Helpers

func eq(t *testing.T, name string, a, b interface{}) bool {
	if a != b {
		t.Error(name, ": (actual)", a, " != ", b, "(expected)")
		return false
	}
	return true
}
<commit_msg>Add some more router tests.<commit_after>package rev

import (
	"fmt"
	"net/http"
	"net/url"
	"regexp"
	"testing"
)

// Data-driven tests that check that a given routes-file line translates into
// the expected Route object.
var routeTestCases = map[string]*Route{
	"get / Application.Index": &Route{
		method:        "GET",
		path:          "/",
		action:        "Application.Index",
		pathPattern:   regexp.MustCompile("/$"),
		staticDir:     "",
		args:          []*arg{},
		actionPattern: regexp.MustCompile("Application\\.Index"),
	},

	"post /app/{id} Application.SaveApp": &Route{
		method:      "POST",
		path:        "/app/{id}",
		action:      "Application.SaveApp",
		pathPattern: regexp.MustCompile("/app/(?P<id>[^/]+)$"),
		staticDir:   "",
		args: []*arg{
			{
				name:       "id",
				constraint: regexp.MustCompile("[^/]+"),
			},
		},
		actionPattern: regexp.MustCompile("Application\\.SaveApp"),
	},

	"post /app/{<[0-9]+>id} Application.SaveApp": &Route{
		method:      "POST",
		path:        "/app/{<[0-9]+>id}",
		action:      "Application.SaveApp",
		pathPattern: regexp.MustCompile("/app/(?P<id>[0-9]+)$"),
		staticDir:   "",
		args: []*arg{
			{
				name:       "id",
				constraint: regexp.MustCompile("[0-9]+"),
			},
		},
		actionPattern: regexp.MustCompile("Application\\.SaveApp"),
	},

	"get /app/? Application.List": &Route{
		method:      "GET",
		path:        "/app/?",
		action:      "Application.List",
		pathPattern: regexp.MustCompile("/app/?$"),
		staticDir:   "",
		args: []*arg{},
		actionPattern: regexp.MustCompile("Application\\.List"),
	},

	"GET /public/ staticDir:www": &Route{
		method:        "GET",
		path:          "/public/",
		action:        "staticDir:www",
		pathPattern:   regexp.MustCompile("^/public/(.*)$"),
		staticDir:     "www",
		args:          []*arg{},
		actionPattern: nil,
	},

	"* /apps/{id}/{action} Application.{action}": &Route{
		method:      "*",
		path:        "/apps/{id}/{action}",
		action:      "Application.{action}",
		pathPattern: regexp.MustCompile("/apps/(?P<id>[^/]+)/(?P<action>[^/]+)$"),
		staticDir:   "",
		args: []*arg{
			{
				name:       "id",
				constraint: regexp.MustCompile("[^/]+"),
			},
			{
				name:       "action",
				constraint: regexp.MustCompile("[^/]+"),
			},
		},
		actionPattern: regexp.MustCompile("Application\\.(?P<action>[^/]+)"),
		actionArgs:    []string{"action"},
	},

	"* /{controller}/{action} {controller}.{action}": &Route{
		method:      "*",
		path:        "/{controller}/{action}",
		action:      "{controller}.{action}",
		pathPattern: regexp.MustCompile("/(?P<controller>[^/]+)/(?P<action>[^/]+)$"),
		staticDir:   "",
		args: []*arg{
			{
				name:       "controller",
				constraint: regexp.MustCompile("[^/]+"),
			},
			{
				name:       "action",
				constraint: regexp.MustCompile("[^/]+"),
			},
		},
		actionPattern: regexp.MustCompile("(?P<controller>[^/]+)\\.(?P<action>[^/]+)"),
		actionArgs:    []string{"controller", "action"},
	},
}

// Run the test cases above.
func TestComputeRoute(t *testing.T) {
	for routeLine, expected := range routeTestCases {
		method, path, action, found := parseRouteLine(routeLine)
		if !found {
			t.Error("Failed to parse route line:", routeLine)
			continue
		}
		actual := NewRoute(method, path, action)
		eq(t, "Method", actual.method, expected.method)
		eq(t, "Path", actual.path, expected.path)
		eq(t, "Action", actual.action, expected.action)
		eq(t, "pathPattern", fmt.Sprint(actual.pathPattern), fmt.Sprint(expected.pathPattern))
		eq(t, "staticDir", actual.staticDir, expected.staticDir)
		eq(t, "len(args)", len(actual.args), len(expected.args))
		for i, arg := range actual.args {
			if len(expected.args) <= i {
				break
			}
			eq(t, "arg.name", arg.name, expected.args[i].name)
			eq(t, "arg.constraint", arg.constraint.String(), expected.args[i].constraint.String())
		}
		eq(t, "actionPattern", fmt.Sprint(actual.actionPattern), fmt.Sprint(expected.actionPattern))
		eq(t, "len(actionArgs)", len(actual.actionArgs), len(expected.actionArgs))
		eq(t, "actionArgs", fmt.Sprint(actual.actionArgs), fmt.Sprint(expected.actionArgs))
		if t.Failed() {
			t.Fatal("Failed on route:", routeLine)
		}
	}
}

// Router Tests

const TEST_ROUTES = `
# This is a comment
GET  /                       Application.Index
GET  /app/{id}               Application.Show
POST /app/{id}               Application.Save

GET  /public/                staticDir:www
*    /{controller}/{action} {controller}.{action}
`

var routeMatchTestCases = map[*http.Request]*RouteMatch{
	&http.Request{
		Method: "GET",
		URL:    &url.URL{Path: "/"},
	}: &RouteMatch{
		ControllerName: "Application",
		MethodName:     "Index",
		Params:         map[string]string{},
		StaticFilename: "",
	},

	&http.Request{
		Method: "GET",
		URL:    &url.URL{Path: "/app/123"},
	}: &RouteMatch{
		ControllerName: "Application",
		MethodName:     "ShowApp",
		Params:         map[string]string{"id": "123"},
		StaticFilename: "",
	},

	&http.Request{
		Method: "POST",
		URL:    &url.URL{Path: "/app/123"},
	}: &RouteMatch{
		ControllerName: "Application",
		MethodName:     "SaveApp",
		Params:         map[string]string{"id": "123"},
		StaticFilename: "",
	},

	&http.Request{
		Method: "GET",
		URL:    &url.URL{Path: "/public/style.css"},
	}: &RouteMatch{
		ControllerName: "",
		MethodName:     "",
		Params:         map[string]string{},
		StaticFilename: "www/style.css",
	},

	&http.Request{
		Method: "GET",
		URL:    &url.URL{Path: "/Implicit/Route"},
	}: &RouteMatch{
		ControllerName: "Implicit",
		MethodName:     "Route",
		Params:         map[string]string{"controller": "Implicit", "action": "Route"},
		StaticFilename: "",
	},
}

func TestRouteMatches(t *testing.T) {
	router := NewRouter(TEST_ROUTES)
	for req, expected := range routeMatchTestCases {
		actual := router.Route(req)
		if !eq(t, "Found route", actual != nil, expected != nil) {
			continue
		}
		eq(t, "ControllerName", actual.ControllerName, expected.ControllerName)
		eq(t, "MethodName", actual.ControllerName, expected.ControllerName)
		eq(t, "len(Params)", len(actual.Params), len(expected.Params))
		for key, actualValue := range actual.Params {
			eq(t, "Params", actualValue, expected.Params[key])
		}
		eq(t, "StaticFilename", actual.StaticFilename, expected.StaticFilename)
	}
}

// Reverse Routing

type ReverseRouteArgs struct {
	action string
	args   map[string]string
}

var reverseRoutingTestCases = map[*ReverseRouteArgs]*ActionDefinition{
	&ReverseRouteArgs{
		action: "Application.Index",
		args:   map[string]string{},
	}: &ActionDefinition{
		Url:    "/",
		Method: "GET",
		Star:   false,
		Action: "Application.Index",
	},

	&ReverseRouteArgs{
		action: "Application.Show",
		args:   map[string]string{"id": "123"},
	}: &ActionDefinition{
		Url:    "/app/123",
		Method: "GET",
		Star:   false,
		Action: "Application.Show",
	},

	&ReverseRouteArgs{
		action: "Implicit.Route",
		args:   map[string]string{},
	}: &ActionDefinition{
		Url:    "/Implicit/Route",
		Method: "GET",
		Star:   true,
		Action: "Implicit.Route",
	},

	&ReverseRouteArgs{
		action: "Application.Save",
		args:   map[string]string{"id": "123", "c": "http://continue"},
	}: &ActionDefinition{
		Url:    "/app/123?c=http%3A%2F%2Fcontinue",
		Method: "POST",
		Star:   false,
		Action: "Application.Save",
	},
}

func TestReverseRouting(t *testing.T) {
	router := NewRouter(TEST_ROUTES)
	for routeArgs, expected := range reverseRoutingTestCases {
		actual := router.Reverse(routeArgs.action, routeArgs.args)
		if !eq(t, "Found route", actual != nil, expected != nil) {
			continue
		}
		eq(t, "Url", actual.Url, expected.Url)
		eq(t, "Method", actual.Method, expected.Method)
		eq(t, "Star", actual.Star, expected.Star)
		eq(t, "Action", actual.Action, expected.Action)
	}
}

// Helpers

func eq(t *testing.T, name string, a, b interface{}) bool {
	if a != b {
		t.Error(name, ": (actual)", a, " != ", b, "(expected)")
		return false
	}
	return true
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc.
//
// 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 main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"io/ioutil"
	"strings"

	"github.com/golang/glog"

	"istio.io/mixer/adapters"
	"istio.io/mixer/adapters/factMapper"
)

func main() {
	const (
		// GRPCPort -- default gRPC server port
		GRPCPort = 9091
		// MaxMessageSize -- default gRPC maximum message size
		MaxMessageSize = 1 * 1024 * 1024
		// MaxConcurrentStreams -- default gRPC concurrency factor
		MaxConcurrentStreams = 2
		// CompressedPayload -- default of whether or not to use compressed gRPC payloads
		CompressedPayload = false
	)

	grpcPort := flag.Int("grpcPort", GRPCPort, "Port exposed for Mixer gRPC API")
	maxMessageSize := flag.Uint("maxMessageSize", MaxMessageSize, "Maximum size of individual gRPC messages")
	maxConcurrentStreams := flag.Uint("maxConcurrentStreams", MaxConcurrentStreams, "Maximum supported number of concurrent gRPC streams")
	compressedPayload := flag.Bool("compressedPayload", CompressedPayload, "Whether to compress gRPC messages")
	serverCertFile := flag.String("serverCertFile", "", "The TLS cert file")
	serverKeyFile := flag.String("serverKeyFile", "", "The TLS key file")
	clientCertFiles := flag.String("clientCertFiles", "", "A set of comma-separated client X509 cert files")
	flag.Parse()

	var err error
	var serverCert *tls.Certificate
	var clientCerts *x509.CertPool

	if *serverCertFile != "" && *serverKeyFile != "" {
		sc, err := tls.LoadX509KeyPair(*serverCertFile, *serverKeyFile)
		if err != nil {
			glog.Exitf("Failed to load server certificate and server key: %v", err)
		}
		serverCert = &sc
	}

	if *clientCertFiles != "" {
		clientCerts = x509.NewCertPool()
		for _, clientCertFile := range strings.Split(*clientCertFiles, ",") {
			pem, err := ioutil.ReadFile(clientCertFile)
			if err != nil {
				glog.Exitf("Failed to load client certificate: %v", err)
			}
			clientCerts.AppendCertsFromPEM(pem)
		}
	}

	var adapterMgr *AdapterManager
	if adapterMgr, err = NewAdapterManager(); err != nil {
		glog.Exitf("Unable to initialize adapters: %v", err)
	}

	// TODO: hackily create a fact mapper builder & adapter.
	// This necessarily needs to be discovered & created through normal
	// adapter config goo, but that doesn't exist yet
	rules := make(map[string]string)
	rules["Lab1"] = "Fact1|Fact2"
	builder := adapterMgr.FactConverters["FactMapper"]
	var adapter adapters.Adapter
	adapter, err = builder.NewAdapter(&factMapper.AdapterConfig{Rules: rules})
	if err != nil {
		glog.Exitf("Unable to create fact conversion adapter " + err.Error())
	}
	factConverter := adapter.(adapters.FactConverter)

	apiServerOptions := APIServerOptions{
		Port:                 uint16(*grpcPort),
		MaxMessageSize:       *maxMessageSize,
		MaxConcurrentStreams: *maxConcurrentStreams,
		CompressedPayload:    *compressedPayload,
		ServerCertificate:    serverCert,
		ClientCertificates:   clientCerts,
		Handlers:             NewAPIHandlers(),
		FactConverter:        factConverter,
	}

	glog.Infof("Starting gRPC server on port %v", apiServerOptions.Port)
	var apiServer *APIServer
	if apiServer, err = NewAPIServer(&apiServerOptions); err != nil {
		glog.Exitf("Unable to initialize API server " + err.Error())
	}
	apiServer.Start()
}
<commit_msg>Switch the mixer to use cobra for command-line parsing. (#55)<commit_after>// Copyright 2016 Google Inc.
//
// 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 main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/golang/glog"
	"github.com/spf13/cobra"

	"istio.io/mixer/adapters"
	"istio.io/mixer/adapters/factMapper"
)

type serverArgs struct {
	port                 uint
	maxMessageSize       uint
	maxConcurrentStreams uint
	compressedPayload    bool
	serverCertFile       string
	serverKeyFile        string
	clientCertFiles      string
}

func printBuilders(listName string, builders map[string]adapters.Builder) {
	fmt.Printf("%s\n", listName)
	if len(builders) > 0 {
		for _, b := range builders {
			fmt.Printf("  %s: %s\n", b.Name(), b.Description())
		}
	} else {
		fmt.Printf("  <none>\n")
	}

	fmt.Printf("\n")
}

func listBuilders() error {
	mgr, err := NewAdapterManager()
	if err != nil {
		return fmt.Errorf("Unable to initialize adapters: %v", err)
	}

	printBuilders("Fact Converters", mgr.FactConverters)
	printBuilders("Fact Updaters", mgr.FactUpdaters)
	printBuilders("List Checkers", mgr.ListCheckers)

	return nil
}

func runServer(sa *serverArgs) error {
	var err error
	var serverCert *tls.Certificate
	var clientCerts *x509.CertPool

	if sa.serverCertFile != "" && sa.serverKeyFile != "" {
		sc, err := tls.LoadX509KeyPair(sa.serverCertFile, sa.serverKeyFile)
		if err != nil {
			return fmt.Errorf("Failed to load server certificate and server key: %v", err)
		}
		serverCert = &sc
	}

	if sa.clientCertFiles != "" {
		clientCerts = x509.NewCertPool()
		for _, clientCertFile := range strings.Split(sa.clientCertFiles, ",") {
			pem, err := ioutil.ReadFile(clientCertFile)
			if err != nil {
				return fmt.Errorf("Failed to load client certificate: %v", err)
			}
			clientCerts.AppendCertsFromPEM(pem)
		}
	}

	var adapterMgr *AdapterManager
	if adapterMgr, err = NewAdapterManager(); err != nil {
		return fmt.Errorf("Unable to initialize adapters: %v", err)
	}

	// TODO: hackily create a fact mapper builder & adapter.
	// This necessarily needs to be discovered & created through normal
	// adapter config goo, but that doesn't exist yet
	rules := make(map[string]string)
	rules["Lab1"] = "Fact1|Fact2"
	builder := adapterMgr.FactConverters["FactMapper"]
	var adapter adapters.Adapter
	adapter, err = builder.NewAdapter(&factMapper.AdapterConfig{Rules: rules})
	if err != nil {
		return fmt.Errorf("Unable to create fact conversion adapter " + err.Error())
	}
	factConverter := adapter.(adapters.FactConverter)

	apiServerOptions := APIServerOptions{
		Port:                 uint16(sa.port),
		MaxMessageSize:       sa.maxMessageSize,
		MaxConcurrentStreams: sa.maxConcurrentStreams,
		CompressedPayload:    sa.compressedPayload,
		ServerCertificate:    serverCert,
		ClientCertificates:   clientCerts,
		Handlers:             NewAPIHandlers(),
		FactConverter:        factConverter,
	}

	var apiServer *APIServer
	if apiServer, err = NewAPIServer(&apiServerOptions); err != nil {
		return fmt.Errorf("Unable to initialize API server " + err.Error())
	}
	apiServer.Start()
	return nil
}

// withArgs is like main except that it is parameterized with the
// command-line arguments to use, along with a function to call
// in case of errors. This allows the function to be invoked
// from test code.
func withArgs(args []string, errorf func(format string, a ...interface{})) {

	rootCmd := cobra.Command{
		Use:   "mixer",
		Short: "The Istio mixer provides control plane functionality to the Istio proxy and services",
	}
	rootCmd.SetArgs(args)

	adapterCmd := cobra.Command{
		Use:   "adapter",
		Short: "Diagnostics for the available mixer adapters",
	}
	rootCmd.AddCommand(&adapterCmd)
	rootCmd.Flags().AddGoFlagSet(flag.CommandLine)

	listCmd := cobra.Command{
		Use:   "list",
		Short: "List available adapter builders",
		Run: func(cmd *cobra.Command, args []string) {
			err := listBuilders()
			if err != nil {
				errorf("%v", err)
			}
		},
	}
	adapterCmd.AddCommand(&listCmd)

	sa := &serverArgs{}
	serverCmd := cobra.Command{
		Use:   "server",
		Short: "Starts the mixer as a server",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Printf("Starting gRPC server on port %v", sa.port)

			err := runServer(sa)
			if err != nil {
				errorf("%v", err)
			}
		},
	}
	serverCmd.PersistentFlags().UintVarP(&sa.port, "port", "p", 9091, "TCP port to use for the mixer's gRPC API")
	serverCmd.PersistentFlags().UintVarP(&sa.maxMessageSize, "maxMessageSize", "", 1024*1024, "Maximum size of individual gRPC messages")
	serverCmd.PersistentFlags().UintVarP(&sa.maxConcurrentStreams, "maxConcurrentStreams", "", 32, "Maximum supported number of concurrent gRPC streams")
	serverCmd.PersistentFlags().BoolVarP(&sa.compressedPayload, "compressedPayload", "", false, "Whether to compress gRPC messages")
	serverCmd.PersistentFlags().StringVarP(&sa.serverCertFile, "serverCertFile", "", "", "The TLS cert file")
	serverCmd.PersistentFlags().StringVarP(&sa.serverKeyFile, "serverKeyFile", "", "", "The TLS key file")
	serverCmd.PersistentFlags().StringVarP(&sa.clientCertFiles, "clientCertFiles", "", "", "A set of comma-separated client X509 cert files")
	rootCmd.AddCommand(&serverCmd)

	if err := rootCmd.Execute(); err != nil {
		errorf("%v", err)
	}
}

func main() {
	withArgs(os.Args[1:],
		func(format string, a ...interface{}) {
			glog.Errorf(format+"\n", a...)
			os.Exit(1)
		})
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/deckarep/golang-set"
	log "github.com/sirupsen/logrus"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

const (
	BUF_SIZE     = 100
	FAKE_NEWLINE = "\\\\"
)

var (
	listen = flag.String("l", ":1991", "listen")
)

var (
	hist  = NewHistory(10)
	chans = mapset.NewSet()
	users = mapset.NewSet()
)

func prepareUsersList() string {
	res := ""
	for user := range users.Iter() {
		str, _ := user.(string)
		res += str + FAKE_NEWLINE
	}
	return res + "\n"
}

func isUserExists(user string) bool {
	for u := range users.Iter() {
		str, _ := u.(string)
		if user == str {
			return true
		}
	}
	return false
}

func main() {
	flag.Parse()

	log.Infof("Start server on %s", *listen)
	ln, err := net.Listen("tcp", *listen)
	defer func() {
		if err := ln.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	if err != nil {
		log.Fatal(err)
	}

	var wg sync.WaitGroup
	quit := make(chan int)
	connChan := make(chan net.Conn)
	broad := make(chan string)

	go func() {
		sigc := make(chan os.Signal, 1)
		signal.Notify(sigc,
			syscall.SIGHUP,
			syscall.SIGINT,
			syscall.SIGTERM,
			syscall.SIGQUIT)

		s := <-sigc
		log.Info(s.String() + " captured. Closing...")
		close(quit)
	}()

	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				log.WithError(err).Error("Failed to accept connection")
				continue
			}
			connChan <- conn
		}
	}()

	for {
		select {
		case conn := <-connChan:
			wg.Add(1)
			go NewChatConnection(conn, broad, quit).Handle(&wg)
		case msg := <-broad:
			log.Print(msg)
			hist.Add(msg)
			for c := range chans.Iter() {
				ch, _ := c.(chan string)
				ch <- msg
			}
		case <-quit:
			wg.Wait()
			return
		}
	}
}
<commit_msg>rename flag<commit_after>package main

import (
	"flag"
	"github.com/deckarep/golang-set"
	log "github.com/sirupsen/logrus"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

const (
	BUF_SIZE     = 100
	FAKE_NEWLINE = "\\\\"
)

var (
	listen = flag.String("b", ":1991", "bind")
)

var (
	hist  = NewHistory(10)
	chans = mapset.NewSet()
	users = mapset.NewSet()
)

func prepareUsersList() string {
	res := ""
	for user := range users.Iter() {
		str, _ := user.(string)
		res += str + FAKE_NEWLINE
	}
	return res + "\n"
}

func isUserExists(user string) bool {
	for u := range users.Iter() {
		str, _ := u.(string)
		if user == str {
			return true
		}
	}
	return false
}

func main() {
	flag.Parse()

	log.Infof("Start server on %s", *listen)
	ln, err := net.Listen("tcp", *listen)
	defer func() {
		if err := ln.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	if err != nil {
		log.Fatal(err)
	}

	var wg sync.WaitGroup
	quit := make(chan int)
	connChan := make(chan net.Conn)
	broad := make(chan string)

	go func() {
		sigc := make(chan os.Signal, 1)
		signal.Notify(sigc,
			syscall.SIGHUP,
			syscall.SIGINT,
			syscall.SIGTERM,
			syscall.SIGQUIT)

		s := <-sigc
		log.Info(s.String() + " captured. Closing...")
		close(quit)
	}()

	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				log.WithError(err).Error("Failed to accept connection")
				continue
			}
			connChan <- conn
		}
	}()

	for {
		select {
		case conn := <-connChan:
			wg.Add(1)
			go NewChatConnection(conn, broad, quit).Handle(&wg)
		case msg := <-broad:
			log.Print(msg)
			hist.Add(msg)
			for c := range chans.Iter() {
				ch, _ := c.(chan string)
				ch <- msg
			}
		case <-quit:
			wg.Wait()
			return
		}
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright 2015 DGraph Labs, Inc.
 *
 * 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 main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"runtime"
	"strings"
	"time"

	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"github.com/Sirupsen/logrus"
	"github.com/dgraph-io/dgraph/commit"
	"github.com/dgraph-io/dgraph/gql"
	"github.com/dgraph-io/dgraph/posting"
	"github.com/dgraph-io/dgraph/query"
	"github.com/dgraph-io/dgraph/query/graph"
	"github.com/dgraph-io/dgraph/rdf"
	"github.com/dgraph-io/dgraph/store"
	"github.com/dgraph-io/dgraph/uid"
	"github.com/dgraph-io/dgraph/worker"
	"github.com/dgraph-io/dgraph/x"
)

var glog = x.Log("server")

var postingDir = flag.String("postings", "", "Directory to store posting lists")
var uidDir = flag.String("uids", "", "XID UID posting lists directory")
var mutationDir = flag.String("mutations", "", "Directory to store mutations")
var port = flag.Int("port", 8080, "Port to run server on.")
var numcpu = flag.Int("numCpu", runtime.NumCPU(),
	"Number of cores to be used by the process")
var instanceIdx = flag.Uint64("instanceIdx", 0,
	"serves only entities whose Fingerprint % numInstance == instanceIdx.")
var workers = flag.String("workers", "",
	"Comma separated list of IP addresses of workers")

func addCorsHeaders(w http.ResponseWriter) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers",
		"Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token,"+
			"X-Auth-Token, Cache-Control, X-Requested-With")
	w.Header().Set("Access-Control-Allow-Credentials", "true")
	w.Header().Set("Connection", "close")
}

func mutationHandler(mu *gql.Mutation) error {
	r := strings.NewReader(mu.Set)
	scanner := bufio.NewScanner(r)
	var nquads []rdf.NQuad
	for scanner.Scan() {
		ln := strings.Trim(scanner.Text(), " \t")
		if len(ln) == 0 {
			continue
		}
		nq, err := rdf.Parse(ln)
		if err != nil {
			glog.WithError(err).Error("While parsing RDF.")
			return err
		}
		nquads = append(nquads, nq)
	}

	xidToUid := make(map[string]uint64)
	for _, nq := range nquads {
		if !strings.HasPrefix(nq.Subject, "_uid_:") {
			xidToUid[nq.Subject] = 0
		}
		if len(nq.ObjectId) > 0 && !strings.HasPrefix(nq.ObjectId, "_uid_:") {
			xidToUid[nq.ObjectId] = 0
		}
	}
	if len(xidToUid) > 0 {
		if err := worker.GetOrAssignUidsOverNetwork(&xidToUid); err != nil {
			glog.WithError(err).Error("GetOrAssignUidsOverNetwork")
			return err
		}
	}

	var edges []x.DirectedEdge
	for _, nq := range nquads {
		edge, err := nq.ToEdgeUsing(xidToUid)
		if err != nil {
			glog.WithField("nquad", nq).WithError(err).
				Error("While converting to edge")
			return err
		}
		edges = append(edges, edge)
	}

	left, err := worker.MutateOverNetwork(edges)
	if err != nil {
		return err
	}
	if len(left) > 0 {
		glog.WithField("left", len(left)).Error("Some edges couldn't be applied")
		for _, e := range left {
			glog.WithField("edge", e).Debug("Unable to apply mutation")
		}
		return fmt.Errorf("Unapplied mutations")
	}
	return nil
}

func queryHandler(w http.ResponseWriter, r *http.Request) {
	addCorsHeaders(w)
	if r.Method == "OPTIONS" {
		return
	}
	if r.Method != "POST" {
		x.SetStatus(w, x.E_INVALID_METHOD, "Invalid method")
		return
	}

	var l query.Latency
	l.Start = time.Now()
	defer r.Body.Close()
	q, err := ioutil.ReadAll(r.Body)
	if err != nil || len(q) == 0 {
		x.Err(glog, err).Error("While reading query")
		x.SetStatus(w, x.E_INVALID_REQUEST, "Invalid request encountered.")
		return
	}

	glog.WithField("q", string(q)).Debug("Query received.")
	gq, mu, err := gql.Parse(string(q))
	if err != nil {
		x.Err(glog, err).Error("While parsing query")
		x.SetStatus(w, x.E_INVALID_REQUEST, err.Error())
		return
	}

	// If we have mutations, run them first.
	if mu != nil && len(mu.Set) > 0 {
		if err = mutationHandler(mu); err != nil {
			glog.WithError(err).Error("While handling mutations.")
			x.SetStatus(w, x.E_ERROR, err.Error())
			return
		}
	}

	if gq == nil || (gq.UID == 0 && len(gq.XID) == 0) {
		x.SetStatus(w, x.E_OK, "Done")
		return
	}

	sg, err := query.ToSubGraph(gq)
	if err != nil {
		x.Err(glog, err).Error("While conversion to internal format")
		x.SetStatus(w, x.E_INVALID_REQUEST, err.Error())
		return
	}
	l.Parsing = time.Since(l.Start)
	glog.WithField("q", string(q)).Debug("Query parsed.")

	rch := make(chan error)
	go query.ProcessGraph(sg, rch, time.Minute)
	err = <-rch
	if err != nil {
		x.Err(glog, err).Error("While executing query")
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
	l.Processing = time.Since(l.Start) - l.Parsing
	glog.WithField("q", string(q)).Debug("Graph processed.")
	js, err := sg.ToJson(&l)
	if err != nil {
		x.Err(glog, err).Error("While converting to Json.")
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
	glog.WithFields(logrus.Fields{
		"total":   time.Since(l.Start),
		"parsing": l.Parsing,
		"process": l.Processing,
		"json":    l.Json,
	}).Info("Query Latencies")

	w.Header().Set("Content-Type", "application/json")
	fmt.Fprint(w, string(js))
}

// server is used to implement graph.DGraphServer
type server struct{}

// This method is used to execute the query and return the response to the
// client as a protocol buffer message.
func (s *server) Query(ctx context.Context,
	req *graph.Request) (resp *graph.Response, err error) {
	if len(req.Query) == 0 {
		glog.Error("While reading query")
		return resp, fmt.Errorf("Empty query")
	}

	var l query.Latency
	l.Start = time.Now()
	// TODO(pawan): Refactor query parsing and graph processing code to a common
	// function used by Query and queryHandler
	glog.WithField("q", req.Query).Debug("Query received.")
	gq, _, err := gql.Parse(req.Query)
	if err != nil {
		x.Err(glog, err).Error("While parsing query")
		return resp, err
	}

	sg, err := query.ToSubGraph(gq)
	if err != nil {
		x.Err(glog, err).Error("While conversion to internal format")
		return resp, err
	}
	l.Parsing = time.Since(l.Start)
	glog.WithField("q", req.Query).Debug("Query parsed.")

	rch := make(chan error)
	go query.ProcessGraph(sg, rch, time.Minute)
	err = <-rch
	if err != nil {
		x.Err(glog, err).Error("While executing query")
		return resp, err
	}
	l.Processing = time.Since(l.Start) - l.Parsing
	glog.WithField("q", req.Query).Debug("Graph processed.")

	node, err := sg.ToProtocolBuffer(&l)
	if err != nil {
		x.Err(glog, err).Error("While converting to protocol buffer.")
		return resp, err
	}
	resp.N = node

	gl := new(graph.Latency)
	gl.Parsing, gl.Processing, gl.Pb = l.Parsing.String(), l.Processing.String(),
		l.ProtocolBuffer.String()
	resp.L = gl

	return resp, err
}

// This function register a DGraph grpc server on the address, which is used
// exchanging protocol buffer messages.
func runGrpcServer(address string) {
	ln, err := net.Listen("tcp", address)
	if err != nil {
		glog.Fatalf("While running server for client: %v", err)
		return
	}
	glog.WithField("address", ln.Addr()).Info("Client Worker listening")

	s := grpc.NewServer()
	graph.RegisterDGraphServer(s, &server{})
	if err = s.Serve(ln); err != nil {
		glog.Fatalf("While serving gRpc requests", err)
	}
	return
}

func main() {
	flag.Parse()
	if !flag.Parsed() {
		glog.Fatal("Unable to parse flags")
	}
	logrus.SetLevel(logrus.InfoLevel)
	numCpus := *numcpu
	prev := runtime.GOMAXPROCS(numCpus)
	glog.WithField("num_cpu", numCpus).WithField("prev_maxprocs", prev).
		Info("Set max procs to num cpus")
	if *port%2 != 0 {
		glog.Fatalf("Port should be an even number: %v", *port)
	}

	ps := new(store.Store)
	ps.Init(*postingDir)
	defer ps.Close()

	clog := commit.NewLogger(*mutationDir, "dgraph", 50<<20)
	clog.SyncEvery = 1
	clog.Init()
	defer clog.Close()

	addrs := strings.Split(*workers, ",")
	lenAddr := uint64(len(addrs))
	if lenAddr == 0 {
		// If no worker is specified, then we're it.
		lenAddr = 1
	}

	posting.Init(clog)
	if *instanceIdx != 0 {
		worker.Init(ps, nil, *instanceIdx, lenAddr)
		uid.Init(nil)
	} else {
		uidStore := new(store.Store)
		uidStore.Init(*uidDir)
		defer uidStore.Close()
		// Only server instance 0 will have uidStore
		worker.Init(ps, uidStore, *instanceIdx, lenAddr)
		uid.Init(uidStore)
	}

	worker.Connect(addrs)
	// Grpc server runs on (port + 1)
	go runGrpcServer(fmt.Sprintf(":%d", *port+1))

	http.HandleFunc("/query", queryHandler)
	glog.WithField("port", *port).Info("Listening for requests...")
	if err := http.ListenAndServe(fmt.Sprintf(":%d", *port), nil); err != nil {
		x.Err(glog, err).Fatal("ListenAndServe")
	}
}
<commit_msg>fixing nil pointer reference bug in server/main.go<commit_after>/*
 * Copyright 2015 DGraph Labs, Inc.
 *
 * 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 main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"runtime"
	"strings"
	"time"

	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"github.com/Sirupsen/logrus"
	"github.com/dgraph-io/dgraph/commit"
	"github.com/dgraph-io/dgraph/gql"
	"github.com/dgraph-io/dgraph/posting"
	"github.com/dgraph-io/dgraph/query"
	"github.com/dgraph-io/dgraph/query/graph"
	"github.com/dgraph-io/dgraph/rdf"
	"github.com/dgraph-io/dgraph/store"
	"github.com/dgraph-io/dgraph/uid"
	"github.com/dgraph-io/dgraph/worker"
	"github.com/dgraph-io/dgraph/x"
)

var glog = x.Log("server")

var postingDir = flag.String("postings", "", "Directory to store posting lists")
var uidDir = flag.String("uids", "", "XID UID posting lists directory")
var mutationDir = flag.String("mutations", "", "Directory to store mutations")
var port = flag.Int("port", 8080, "Port to run server on.")
var numcpu = flag.Int("numCpu", runtime.NumCPU(),
	"Number of cores to be used by the process")
var instanceIdx = flag.Uint64("instanceIdx", 0,
	"serves only entities whose Fingerprint % numInstance == instanceIdx.")
var workers = flag.String("workers", "",
	"Comma separated list of IP addresses of workers")

func addCorsHeaders(w http.ResponseWriter) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers",
		"Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token,"+
			"X-Auth-Token, Cache-Control, X-Requested-With")
	w.Header().Set("Access-Control-Allow-Credentials", "true")
	w.Header().Set("Connection", "close")
}

func mutationHandler(mu *gql.Mutation) error {
	r := strings.NewReader(mu.Set)
	scanner := bufio.NewScanner(r)
	var nquads []rdf.NQuad
	for scanner.Scan() {
		ln := strings.Trim(scanner.Text(), " \t")
		if len(ln) == 0 {
			continue
		}
		nq, err := rdf.Parse(ln)
		if err != nil {
			glog.WithError(err).Error("While parsing RDF.")
			return err
		}
		nquads = append(nquads, nq)
	}

	xidToUid := make(map[string]uint64)
	for _, nq := range nquads {
		if !strings.HasPrefix(nq.Subject, "_uid_:") {
			xidToUid[nq.Subject] = 0
		}
		if len(nq.ObjectId) > 0 && !strings.HasPrefix(nq.ObjectId, "_uid_:") {
			xidToUid[nq.ObjectId] = 0
		}
	}
	if len(xidToUid) > 0 {
		if err := worker.GetOrAssignUidsOverNetwork(&xidToUid); err != nil {
			glog.WithError(err).Error("GetOrAssignUidsOverNetwork")
			return err
		}
	}

	var edges []x.DirectedEdge
	for _, nq := range nquads {
		edge, err := nq.ToEdgeUsing(xidToUid)
		if err != nil {
			glog.WithField("nquad", nq).WithError(err).
				Error("While converting to edge")
			return err
		}
		edges = append(edges, edge)
	}

	left, err := worker.MutateOverNetwork(edges)
	if err != nil {
		return err
	}
	if len(left) > 0 {
		glog.WithField("left", len(left)).Error("Some edges couldn't be applied")
		for _, e := range left {
			glog.WithField("edge", e).Debug("Unable to apply mutation")
		}
		return fmt.Errorf("Unapplied mutations")
	}
	return nil
}

func queryHandler(w http.ResponseWriter, r *http.Request) {
	addCorsHeaders(w)
	if r.Method == "OPTIONS" {
		return
	}
	if r.Method != "POST" {
		x.SetStatus(w, x.E_INVALID_METHOD, "Invalid method")
		return
	}

	var l query.Latency
	l.Start = time.Now()
	defer r.Body.Close()
	q, err := ioutil.ReadAll(r.Body)
	if err != nil || len(q) == 0 {
		x.Err(glog, err).Error("While reading query")
		x.SetStatus(w, x.E_INVALID_REQUEST, "Invalid request encountered.")
		return
	}

	glog.WithField("q", string(q)).Debug("Query received.")
	gq, mu, err := gql.Parse(string(q))
	if err != nil {
		x.Err(glog, err).Error("While parsing query")
		x.SetStatus(w, x.E_INVALID_REQUEST, err.Error())
		return
	}

	// If we have mutations, run them first.
	if mu != nil && len(mu.Set) > 0 {
		if err = mutationHandler(mu); err != nil {
			glog.WithError(err).Error("While handling mutations.")
			x.SetStatus(w, x.E_ERROR, err.Error())
			return
		}
	}

	if gq == nil || (gq.UID == 0 && len(gq.XID) == 0) {
		x.SetStatus(w, x.E_OK, "Done")
		return
	}

	sg, err := query.ToSubGraph(gq)
	if err != nil {
		x.Err(glog, err).Error("While conversion to internal format")
		x.SetStatus(w, x.E_INVALID_REQUEST, err.Error())
		return
	}
	l.Parsing = time.Since(l.Start)
	glog.WithField("q", string(q)).Debug("Query parsed.")

	rch := make(chan error)
	go query.ProcessGraph(sg, rch, time.Minute)
	err = <-rch
	if err != nil {
		x.Err(glog, err).Error("While executing query")
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
	l.Processing = time.Since(l.Start) - l.Parsing
	glog.WithField("q", string(q)).Debug("Graph processed.")
	js, err := sg.ToJson(&l)
	if err != nil {
		x.Err(glog, err).Error("While converting to Json.")
		x.SetStatus(w, x.E_ERROR, err.Error())
		return
	}
	glog.WithFields(logrus.Fields{
		"total":   time.Since(l.Start),
		"parsing": l.Parsing,
		"process": l.Processing,
		"json":    l.Json,
	}).Info("Query Latencies")

	w.Header().Set("Content-Type", "application/json")
	fmt.Fprint(w, string(js))
}

// server is used to implement graph.DGraphServer
type server struct{}

// This method is used to execute the query and return the response to the
// client as a protocol buffer message.
func (s *server) Query(ctx context.Context,
	req *graph.Request) (*graph.Response, error) {
	resp := new(graph.Response)
	if len(req.Query) == 0 {
		glog.Error("While reading query")
		return resp, fmt.Errorf("Empty query")
	}

	var l query.Latency
	l.Start = time.Now()
	// TODO(pawan): Refactor query parsing and graph processing code to a common
	// function used by Query and queryHandler
	glog.WithField("q", req.Query).Debug("Query received.")
	gq, _, err := gql.Parse(req.Query)
	if err != nil {
		x.Err(glog, err).Error("While parsing query")
		return resp, err
	}

	sg, err := query.ToSubGraph(gq)
	if err != nil {
		x.Err(glog, err).Error("While conversion to internal format")
		return resp, err
	}
	l.Parsing = time.Since(l.Start)
	glog.WithField("q", req.Query).Debug("Query parsed.")

	rch := make(chan error)
	go query.ProcessGraph(sg, rch, time.Minute)
	err = <-rch
	if err != nil {
		x.Err(glog, err).Error("While executing query")
		return resp, err
	}
	l.Processing = time.Since(l.Start) - l.Parsing
	glog.WithField("q", req.Query).Debug("Graph processed.")

	node, err := sg.ToProtocolBuffer(&l)
	if err != nil {
		x.Err(glog, err).Error("While converting to protocol buffer.")
		return resp, err
	}
	resp.N = node

	gl := new(graph.Latency)
	gl.Parsing, gl.Processing, gl.Pb = l.Parsing.String(), l.Processing.String(),
		l.ProtocolBuffer.String()
	resp.L = gl

	return resp, err
}

// This function register a DGraph grpc server on the address, which is used
// exchanging protocol buffer messages.
func runGrpcServer(address string) {
	ln, err := net.Listen("tcp", address)
	if err != nil {
		glog.Fatalf("While running server for client: %v", err)
		return
	}
	glog.WithField("address", ln.Addr()).Info("Client Worker listening")

	s := grpc.NewServer()
	graph.RegisterDGraphServer(s, &server{})
	if err = s.Serve(ln); err != nil {
		glog.Fatalf("While serving gRpc requests", err)
	}
	return
}

func main() {
	flag.Parse()
	if !flag.Parsed() {
		glog.Fatal("Unable to parse flags")
	}
	logrus.SetLevel(logrus.InfoLevel)
	numCpus := *numcpu
	prev := runtime.GOMAXPROCS(numCpus)
	glog.WithField("num_cpu", numCpus).WithField("prev_maxprocs", prev).
		Info("Set max procs to num cpus")
	if *port%2 != 0 {
		glog.Fatalf("Port should be an even number: %v", *port)
	}

	ps := new(store.Store)
	ps.Init(*postingDir)
	defer ps.Close()

	clog := commit.NewLogger(*mutationDir, "dgraph", 50<<20)
	clog.SyncEvery = 1
	clog.Init()
	defer clog.Close()

	addrs := strings.Split(*workers, ",")
	lenAddr := uint64(len(addrs))
	if lenAddr == 0 {
		// If no worker is specified, then we're it.
		lenAddr = 1
	}

	posting.Init(clog)
	if *instanceIdx != 0 {
		worker.Init(ps, nil, *instanceIdx, lenAddr)
		uid.Init(nil)
	} else {
		uidStore := new(store.Store)
		uidStore.Init(*uidDir)
		defer uidStore.Close()
		// Only server instance 0 will have uidStore
		worker.Init(ps, uidStore, *instanceIdx, lenAddr)
		uid.Init(uidStore)
	}

	worker.Connect(addrs)
	// Grpc server runs on (port + 1)
	go runGrpcServer(fmt.Sprintf(":%d", *port+1))

	http.HandleFunc("/query", queryHandler)
	glog.WithField("port", *port).Info("Listening for requests...")
	if err := http.ListenAndServe(fmt.Sprintf(":%d", *port), nil); err != nil {
		x.Err(glog, err).Fatal("ListenAndServe")
	}
}
<|endoftext|>


<commit_before>package kcp

import (
	"crypto/md5"
	"crypto/rand"
	"io"
	"sync"
)

// nonceMD5 is a nonce generator for each packet header
// which took the advantages of both MD5 and CSPRNG(like /dev/urandom).
// The benchmark shows it's faster than previous CSPRNG only method.
type nonceMD5 struct {
	data [md5.Size]byte
	mu   sync.Mutex
}

// Nonce fills a nonce into the provided slice with no more than md5.Size bytes
// the entropy will be updated whenenver a leading 0 is found
func (n *nonceMD5) Fill(nonce []byte) {
	n.mu.Lock()
	if n.data[0] == 0 { // 1/256 chance for entropy update
		io.ReadFull(rand.Reader, n.data[:])
	}
	n.data = md5.Sum(n.data[:])
	copy(nonce, n.data[:])
	n.mu.Unlock()
	return
}
<commit_msg>remove the lock in rand<commit_after>package kcp

import (
	"crypto/md5"
	"crypto/rand"
	"io"
)

// nonceMD5 is a nonce generator for each packet header
// which took the advantages of both MD5 and CSPRNG(like /dev/urandom).
// The benchmark shows it's faster than previous CSPRNG only method.
type nonceMD5 struct {
	data [md5.Size]byte
}

// Nonce fills a nonce into the provided slice with no more than md5.Size bytes
// the entropy will be updated whenever a leading 0 appears
func (n *nonceMD5) Fill(nonce []byte) {
	if n.data[0] == 0 { // 1/256 chance for entropy update
		io.ReadFull(rand.Reader, n.data[:])
	}
	n.data = md5.Sum(n.data[:])
	copy(nonce, n.data[:])
	return
}
<|endoftext|>


<commit_before>package udpsender

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha1"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"time"
)

type RecvService struct {
	key          string
	ch           chan *RecvBlock
	transactions map[string]*RecvTransaction
	sock         *net.UDPConn
	fn           RecvHandleFunc
}

type RecvTransaction struct {
	num         int
	blocks      map[int]*RecvBlock
	lastTime    int64
	blockSize   int64
	Speed       float64
	completed   bool
	Addr        *net.UDPAddr
	Transaction []byte
}

type RecvBlock struct {
	transaction      []byte
	blockHash        []byte
	fileHash         []byte
	fileData         []byte
	num, id          int32
	start, end, size int64
	time1            int64
	addr             *net.UDPAddr
}

type RecvHandleFunc func(service *RecvService, transaction *RecvTransaction)

func InitRecvService(key string, url string, fn RecvHandleFunc) (service *RecvService, err error) {
	service = new(RecvService)
	service.key = key
	service.fn = fn
	service.ch = make(chan *RecvBlock, 100)
	service.transactions = map[string]*RecvTransaction{}

	addr, err := net.ResolveUDPAddr("udp", url)
	if err != nil {
		return nil, err
	}
	service.sock, err = net.ListenUDP("udp", addr)
	if err != nil {
		return nil, err
	}
	go service.loop()

	return service, nil
}

func InitRecvTransaction(block *RecvBlock) *RecvTransaction {
	transaction := new(RecvTransaction)
	transaction.blocks = map[int]*RecvBlock{}
	transaction.lastTime = time.Now().UnixNano()
	transaction.blockSize = int64(len(block.fileData))
	transaction.num = int(block.num)
	transaction.Addr = block.addr
	transaction.Transaction = block.transaction
	return transaction
}
func (service *RecvService) processRecvBlock(block *RecvBlock) {
	stransaction := hex.EncodeToString(block.transaction)
	transaction, ok := service.transactions[stransaction]
	if !ok {
		fmt.Println("Creating transaction:", stransaction)
		transaction = InitRecvTransaction(block)
		service.transactions[stransaction] = transaction
	}
	if transaction.num != int(block.num) {
		// error
	}
	block.time1 = time.Now().UnixNano()
	fdtime := float64(block.time1 - transaction.lastTime)
	transaction.lastTime = block.time1
	speed := float64(len(block.fileData)) / fdtime * float64(time.Second)
	transaction.Speed = movingExpAvg(speed, transaction.Speed, fdtime, float64(time.Second))

	_, ok = transaction.blocks[int(block.id)]
	if ok {
		// dupe
		return
	}
	transaction.blocks[int(block.id)] = block
	// fmt.Printf("got %d %d %d\n%s\n", block.id, len(transaction.blocks), transaction.num, string(block.fileData))
	fmt.Printf("got %d %d %d %f\n", block.id, len(transaction.blocks), transaction.num, transaction.Speed)
	if len(transaction.blocks) == transaction.num {
		transaction.completed = true
		service.fn(service, transaction)
		// fmt.Printf("Finish RecvTransaction %f\n", transaction.Speed)
		// // finish transaction
		// // save file

		// f1, err := os.Create("incomming/" + stransaction + "/file.out")
		// if err != nil {
		// 	panic(err)
		// }
		// // defer f1.Close()

		// for i := 1; i <= transaction.num; i++ {
		// 	block1, ok := transaction.blocks[i]
		// 	if !ok {
		// 		// error
		// 		panic(errors.New(fmt.Sprintf("Not all blocks %d", i)))
		// 	}
		// 	f1.Write(block1.fileData)
		// }
		// f1.Close()
		// file recieved
	}
}
func (service *RecvService) WriteTo(w io.Writer) error {
	for i := 1; i <= transaction.num; i++ {
		block1, ok := transaction.blocks[i]
		if !ok {
			// error
			return errors.New(fmt.Sprintf("Not all blocks %d", i))
		}
		w.Write(block1.fileData)
	}
	return nil
}
func (service *RecvService) Loop() {

	for {
		// len, remote, err := sock.ReadFromUDP(buf[:])
		var buf [1 << 16]byte
		len, remote, err := service.sock.ReadFromUDP(buf[:])
		if err != nil {
			panic(err)
		}
		service.processPacket(remote, buf[:len])
	}
}
func (service *RecvService) loop() {
	ticker := time.Tick(100 * time.Millisecond)
	for {

		select {
		case blck, ok := <-service.ch:
			if ok {
				service.processRecvBlock(blck)
			}
		case <-ticker:
			time1 := time.Now().UnixNano()
			for _, transaction := range service.transactions {
				if transaction.completed {
					continue
				}

				fdtime := float64(time1 - transaction.lastTime)
				transaction.lastTime = time1
				transaction.Speed = movingExpAvg(0.0, transaction.Speed, fdtime, float64(time.Second))

				//
				// find last_block
				last_block := -1
				var block *RecvBlock
				for i, _block := range transaction.blocks {
					if last_block < i {
						last_block = i
						block = _block
					}
				}
				if last_block == -1 {
					continue
				}
				// dtime  = Now - last_block.time
				dtime := time1 - block.time1
				// last_block + (dtime-100ms)*speed/blocksize
				shift := last_block + int((float64(dtime)-100.0*float64(time.Millisecond))*transaction.Speed/float64(transaction.blockSize))
				// get first 100 un recieved blocks

				req := []int{}

				for i := 1; i <= transaction.num; i++ {
					if i > shift {
						break
					}
					_, ok := transaction.blocks[i]
					if !ok {
						req = append(req, i)
						if len(req) >= 100 {
							break
						}
					}
				}

				if len(req) > 0 {
					service.rerequestPackets(transaction, req)
				}
			}
		}

	}

}
func (service *RecvService) rerequestPackets(transaction *RecvTransaction, req []int) {
	r1 := make([]byte, 16)
	_, err := rand.Read(r1)
	if err != nil {
		panic(err)
	}

	data := append(transaction.Transaction, r1...)

	buffer := new(bytes.Buffer)
	binary.Write(buffer, binary.LittleEndian, int32(1))
	binary.Write(buffer, binary.LittleEndian, int32(len(req)))
	for _, i := range req {
		binary.Write(buffer, binary.LittleEndian, int32(i))
	}

	block, err := aes.NewCipher([]byte(service.key))
	if err != nil {
		panic(err)
	}

	data1 := buffer.Bytes()

	stream := cipher.NewCFBEncrypter(block, r1)
	stream.XORKeyStream(data1, data1)

	data = append(data, data1...)

	service.sock.WriteToUDP(data, transaction.Addr)
}
func (service *RecvService) processPacket(addr *net.UDPAddr, buf []byte) {
	blck := new(RecvBlock)
	blck.transaction = buf[:16]
	blck.addr = addr
	r1 := buf[16:32]
	data2 := buf[32:]

	block, err := aes.NewCipher([]byte(service.key))
	if err != nil {
		panic(err)
	}

	stream := cipher.NewCFBDecrypter(block, r1)

	stream.XORKeyStream(data2, data2)

	// fileHash := data2[0:20]

	buffer := bytes.NewBuffer(data2[20:52])
	binary.Read(buffer, binary.LittleEndian, &blck.id)
	binary.Read(buffer, binary.LittleEndian, &blck.num)
	binary.Read(buffer, binary.LittleEndian, &blck.start)
	binary.Read(buffer, binary.LittleEndian, &blck.end)
	binary.Read(buffer, binary.LittleEndian, &blck.size)
	blck.blockHash = data2[52:72]
	bufCopy := make([]byte, len(data2[72:]))
	copy(bufCopy[:], data2[72:])
	blck.fileData = bufCopy
	blockHash := sha1.Sum(blck.fileData)
	verified := bytes.Equal(blck.blockHash, blockHash[:])

	fmt.Println("processPacker:", blck.id, hex.EncodeToString(r1))

	if verified {
		// os.MkdirAll("incomming/"+hex.EncodeToString(blck.transaction), 0700)
		// f1, err := os.Create("incomming/" + hex.EncodeToString(blck.transaction) + "/" + fmt.Sprintf("%d.data", blck.id))
		// if err != nil {
		// 	panic(err)
		// }
		// defer f1.Close()
		// f1.Write(blck.fileData)
		service.ch <- blck
	} else {
		fmt.Println("wrong hash %#v %#v", blck.blockHash, blockHash[:])
	}

}
<commit_msg>fix WriteTo<commit_after>package udpsender

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha1"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"time"
)

type RecvService struct {
	key          string
	ch           chan *RecvBlock
	transactions map[string]*RecvTransaction
	sock         *net.UDPConn
	fn           RecvHandleFunc
}

type RecvTransaction struct {
	num         int
	blocks      map[int]*RecvBlock
	lastTime    int64
	blockSize   int64
	Speed       float64
	completed   bool
	Addr        *net.UDPAddr
	Transaction []byte
}

type RecvBlock struct {
	transaction      []byte
	blockHash        []byte
	fileHash         []byte
	fileData         []byte
	num, id          int32
	start, end, size int64
	time1            int64
	addr             *net.UDPAddr
}

type RecvHandleFunc func(service *RecvService, transaction *RecvTransaction)

func InitRecvService(key string, url string, fn RecvHandleFunc) (service *RecvService, err error) {
	service = new(RecvService)
	service.key = key
	service.fn = fn
	service.ch = make(chan *RecvBlock, 100)
	service.transactions = map[string]*RecvTransaction{}

	addr, err := net.ResolveUDPAddr("udp", url)
	if err != nil {
		return nil, err
	}
	service.sock, err = net.ListenUDP("udp", addr)
	if err != nil {
		return nil, err
	}
	go service.loop()

	return service, nil
}

func InitRecvTransaction(block *RecvBlock) *RecvTransaction {
	transaction := new(RecvTransaction)
	transaction.blocks = map[int]*RecvBlock{}
	transaction.lastTime = time.Now().UnixNano()
	transaction.blockSize = int64(len(block.fileData))
	transaction.num = int(block.num)
	transaction.Addr = block.addr
	transaction.Transaction = block.transaction
	return transaction
}
func (service *RecvService) processRecvBlock(block *RecvBlock) {
	stransaction := hex.EncodeToString(block.transaction)
	transaction, ok := service.transactions[stransaction]
	if !ok {
		fmt.Println("Creating transaction:", stransaction)
		transaction = InitRecvTransaction(block)
		service.transactions[stransaction] = transaction
	}
	if transaction.num != int(block.num) {
		// error
	}
	block.time1 = time.Now().UnixNano()
	fdtime := float64(block.time1 - transaction.lastTime)
	transaction.lastTime = block.time1
	speed := float64(len(block.fileData)) / fdtime * float64(time.Second)
	transaction.Speed = movingExpAvg(speed, transaction.Speed, fdtime, float64(time.Second))

	_, ok = transaction.blocks[int(block.id)]
	if ok {
		// dupe
		return
	}
	transaction.blocks[int(block.id)] = block
	// fmt.Printf("got %d %d %d\n%s\n", block.id, len(transaction.blocks), transaction.num, string(block.fileData))
	fmt.Printf("got %d %d %d %f\n", block.id, len(transaction.blocks), transaction.num, transaction.Speed)
	if len(transaction.blocks) == transaction.num {
		transaction.completed = true
		service.fn(service, transaction)
		// fmt.Printf("Finish RecvTransaction %f\n", transaction.Speed)
		// // finish transaction
		// // save file

		// f1, err := os.Create("incomming/" + stransaction + "/file.out")
		// if err != nil {
		// 	panic(err)
		// }
		// // defer f1.Close()

		// for i := 1; i <= transaction.num; i++ {
		// 	block1, ok := transaction.blocks[i]
		// 	if !ok {
		// 		// error
		// 		panic(errors.New(fmt.Sprintf("Not all blocks %d", i)))
		// 	}
		// 	f1.Write(block1.fileData)
		// }
		// f1.Close()
		// file recieved
	}
}
func (transaction *RecvTransaction) WriteTo(w io.Writer) error {
	for i := 1; i <= transaction.num; i++ {
		block1, ok := transaction.blocks[i]
		if !ok {
			// error
			return errors.New(fmt.Sprintf("Not all blocks %d", i))
		}
		w.Write(block1.fileData)
	}
	return nil
}
func (service *RecvService) Loop() {

	for {
		// len, remote, err := sock.ReadFromUDP(buf[:])
		var buf [1 << 16]byte
		len, remote, err := service.sock.ReadFromUDP(buf[:])
		if err != nil {
			panic(err)
		}
		service.processPacket(remote, buf[:len])
	}
}
func (service *RecvService) loop() {
	ticker := time.Tick(100 * time.Millisecond)
	for {

		select {
		case blck, ok := <-service.ch:
			if ok {
				service.processRecvBlock(blck)
			}
		case <-ticker:
			time1 := time.Now().UnixNano()
			for _, transaction := range service.transactions {
				if transaction.completed {
					continue
				}

				fdtime := float64(time1 - transaction.lastTime)
				transaction.lastTime = time1
				transaction.Speed = movingExpAvg(0.0, transaction.Speed, fdtime, float64(time.Second))

				//
				// find last_block
				last_block := -1
				var block *RecvBlock
				for i, _block := range transaction.blocks {
					if last_block < i {
						last_block = i
						block = _block
					}
				}
				if last_block == -1 {
					continue
				}
				// dtime  = Now - last_block.time
				dtime := time1 - block.time1
				// last_block + (dtime-100ms)*speed/blocksize
				shift := last_block + int((float64(dtime)-100.0*float64(time.Millisecond))*transaction.Speed/float64(transaction.blockSize))
				// get first 100 un recieved blocks

				req := []int{}

				for i := 1; i <= transaction.num; i++ {
					if i > shift {
						break
					}
					_, ok := transaction.blocks[i]
					if !ok {
						req = append(req, i)
						if len(req) >= 100 {
							break
						}
					}
				}

				if len(req) > 0 {
					service.rerequestPackets(transaction, req)
				}
			}
		}

	}

}
func (service *RecvService) rerequestPackets(transaction *RecvTransaction, req []int) {
	r1 := make([]byte, 16)
	_, err := rand.Read(r1)
	if err != nil {
		panic(err)
	}

	data := append(transaction.Transaction, r1...)

	buffer := new(bytes.Buffer)
	binary.Write(buffer, binary.LittleEndian, int32(1))
	binary.Write(buffer, binary.LittleEndian, int32(len(req)))
	for _, i := range req {
		binary.Write(buffer, binary.LittleEndian, int32(i))
	}

	block, err := aes.NewCipher([]byte(service.key))
	if err != nil {
		panic(err)
	}

	data1 := buffer.Bytes()

	stream := cipher.NewCFBEncrypter(block, r1)
	stream.XORKeyStream(data1, data1)

	data = append(data, data1...)

	service.sock.WriteToUDP(data, transaction.Addr)
}
func (service *RecvService) processPacket(addr *net.UDPAddr, buf []byte) {
	blck := new(RecvBlock)
	blck.transaction = buf[:16]
	blck.addr = addr
	r1 := buf[16:32]
	data2 := buf[32:]

	block, err := aes.NewCipher([]byte(service.key))
	if err != nil {
		panic(err)
	}

	stream := cipher.NewCFBDecrypter(block, r1)

	stream.XORKeyStream(data2, data2)

	// fileHash := data2[0:20]

	buffer := bytes.NewBuffer(data2[20:52])
	binary.Read(buffer, binary.LittleEndian, &blck.id)
	binary.Read(buffer, binary.LittleEndian, &blck.num)
	binary.Read(buffer, binary.LittleEndian, &blck.start)
	binary.Read(buffer, binary.LittleEndian, &blck.end)
	binary.Read(buffer, binary.LittleEndian, &blck.size)
	blck.blockHash = data2[52:72]
	bufCopy := make([]byte, len(data2[72:]))
	copy(bufCopy[:], data2[72:])
	blck.fileData = bufCopy
	blockHash := sha1.Sum(blck.fileData)
	verified := bytes.Equal(blck.blockHash, blockHash[:])

	fmt.Println("processPacker:", blck.id, hex.EncodeToString(r1))

	if verified {
		// os.MkdirAll("incomming/"+hex.EncodeToString(blck.transaction), 0700)
		// f1, err := os.Create("incomming/" + hex.EncodeToString(blck.transaction) + "/" + fmt.Sprintf("%d.data", blck.id))
		// if err != nil {
		// 	panic(err)
		// }
		// defer f1.Close()
		// f1.Write(blck.fileData)
		service.ch <- blck
	} else {
		fmt.Println("wrong hash %#v %#v", blck.blockHash, blockHash[:])
	}

}
<|endoftext|>