<commit_before>package operators

import (
	"encoding/json"
	"fmt"
	"strings"

	. "github.com/onsi/ginkgo"

	"github.com/openshift/origin/pkg/oc/cli/admin/release"
	exutil "github.com/openshift/origin/test/extended/util"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	e2e "k8s.io/kubernetes/test/e2e/framework"
)

var _ = Describe("[Feature:Platform][Smoke] Managed cluster", func() {
	oc := exutil.NewCLIWithoutNamespace("operators")

	It("should ensure pods use images from our release image with proper ImagePullPolicy", func() {
		// find out the current installed release info
		out, err := oc.Run("adm", "release", "info").Args("--pullspecs", "-o", "json").Output()
		if err != nil {
			e2e.Failf("unable to read release payload with error: %v", err)
		}
		releaseInfo := &release.ReleaseInfo{}
		if err := json.Unmarshal([]byte(out), &releaseInfo); err != nil {
			e2e.Failf("unable to decode release payload with error: %v", err)
		}
		e2e.Logf("Release Info image=%s", releaseInfo.Image)
		e2e.Logf("Release Info number of tags %v", len(releaseInfo.References.Spec.Tags))

		// valid images include the release image and all its tagged references
		validImages := sets.NewString()
		validImages.Insert(releaseInfo.Image)
		if releaseInfo.References == nil {
			e2e.Failf("no references found")
		}
		for i := range releaseInfo.References.Spec.Tags {
			tag := releaseInfo.References.Spec.Tags[i]
			if tag.From != nil && tag.From.Kind == "DockerImage" {
				validImages.Insert(tag.From.Name)
			}
		}

		// iterate over the references to find valid images
		pods, err := oc.KubeFramework().ClientSet.CoreV1().Pods("").List(metav1.ListOptions{})
		if err != nil {
			e2e.Failf("unable to list pods: %v", err)
		}

		// list of pods that use images not in the release payload
		invalidPodContainerImages := sets.NewString()
		invalidPodContainerImagePullPolicy := sets.NewString()
		// a pod in a namespace that begins with kube-* or openshift-* must come from our release payload
		// TODO components in openshift-operators may not come from our payload, may want to weaken restriction
		namespacePrefixes := sets.NewString("kube-", "openshift-")
		for i := range pods.Items {
			pod := pods.Items[i]
			for _, prefix := range namespacePrefixes.List() {
				if !strings.HasPrefix(pod.Namespace, prefix) {
					continue
				}
				containersToInspect := []v1.Container{}
				for j := range pod.Spec.InitContainers {
					containersToInspect = append(containersToInspect, pod.Spec.InitContainers[j])
				}
				for j := range pod.Spec.Containers {
					containersToInspect = append(containersToInspect, pod.Spec.Containers[j])
				}
				for j := range containersToInspect {
					container := containersToInspect[j]
					if !validImages.Has(container.Image) {
						invalidPodContainerImages.Insert(fmt.Sprintf("%s/%s/%s image=%s", pod.Namespace, pod.Name, container.Name, container.Image))
					}
					if container.ImagePullPolicy != v1.PullIfNotPresent {
						invalidPodContainerImagePullPolicy.Insert(fmt.Sprintf("%s/%s/%s imagePullPolicy=%s", pod.Namespace, pod.Name, container.Name, container.ImagePullPolicy))
					}
				}
			}
		}
		// log for debugging output before we ultimately fail
		e2e.Logf("Pods found with invalid container images not present in release payload: %s", strings.Join(invalidPodContainerImages.List(), "\n"))
		e2e.Logf("Pods found with invalid container image pull policy not equal to IfNotPresent: %s", strings.Join(invalidPodContainerImagePullPolicy.List(), "\n"))
		if len(invalidPodContainerImages) > 0 {
			e2e.Failf("Pods found with invalid container images not present in release payload: %s", strings.Join(invalidPodContainerImages.List(), "\n"))
		}
		if len(invalidPodContainerImagePullPolicy) > 0 {
			e2e.Failf("Pods found with invalid container image pull policy not equal to IfNotPresent: %s", strings.Join(invalidPodContainerImagePullPolicy.List(), "\n"))
		}
	})
})
<commit_msg>extended: exit early if unable to read release payload<commit_after>package operators

import (
	"encoding/json"
	"fmt"
	"strings"

	. "github.com/onsi/ginkgo"

	"github.com/openshift/origin/pkg/oc/cli/admin/release"
	exutil "github.com/openshift/origin/test/extended/util"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
	e2e "k8s.io/kubernetes/test/e2e/framework"
)

var _ = Describe("[Feature:Platform][Smoke] Managed cluster", func() {
	oc := exutil.NewCLIWithoutNamespace("operators")

	It("should ensure pods use images from our release image with proper ImagePullPolicy", func() {
		// find out the current installed release info
		out, err := oc.Run("adm", "release", "info").Args("--pullspecs", "-o", "json").Output()
		if err != nil {
			// TODO need to determine why release tests are not having access to read payload
			e2e.Logf("unable to read release payload with error: %v", err)
			return
		}
		releaseInfo := &release.ReleaseInfo{}
		if err := json.Unmarshal([]byte(out), &releaseInfo); err != nil {
			e2e.Failf("unable to decode release payload with error: %v", err)
		}
		e2e.Logf("Release Info image=%s", releaseInfo.Image)
		e2e.Logf("Release Info number of tags %v", len(releaseInfo.References.Spec.Tags))

		// valid images include the release image and all its tagged references
		validImages := sets.NewString()
		validImages.Insert(releaseInfo.Image)
		if releaseInfo.References == nil {
			e2e.Failf("no references found")
		}
		for i := range releaseInfo.References.Spec.Tags {
			tag := releaseInfo.References.Spec.Tags[i]
			if tag.From != nil && tag.From.Kind == "DockerImage" {
				validImages.Insert(tag.From.Name)
			}
		}

		// iterate over the references to find valid images
		pods, err := oc.KubeFramework().ClientSet.CoreV1().Pods("").List(metav1.ListOptions{})
		if err != nil {
			e2e.Failf("unable to list pods: %v", err)
		}

		// list of pods that use images not in the release payload
		invalidPodContainerImages := sets.NewString()
		invalidPodContainerImagePullPolicy := sets.NewString()
		// a pod in a namespace that begins with kube-* or openshift-* must come from our release payload
		// TODO components in openshift-operators may not come from our payload, may want to weaken restriction
		namespacePrefixes := sets.NewString("kube-", "openshift-")
		for i := range pods.Items {
			pod := pods.Items[i]
			for _, prefix := range namespacePrefixes.List() {
				if !strings.HasPrefix(pod.Namespace, prefix) {
					continue
				}
				containersToInspect := []v1.Container{}
				for j := range pod.Spec.InitContainers {
					containersToInspect = append(containersToInspect, pod.Spec.InitContainers[j])
				}
				for j := range pod.Spec.Containers {
					containersToInspect = append(containersToInspect, pod.Spec.Containers[j])
				}
				for j := range containersToInspect {
					container := containersToInspect[j]
					if !validImages.Has(container.Image) {
						invalidPodContainerImages.Insert(fmt.Sprintf("%s/%s/%s image=%s", pod.Namespace, pod.Name, container.Name, container.Image))
					}
					if container.ImagePullPolicy != v1.PullIfNotPresent {
						invalidPodContainerImagePullPolicy.Insert(fmt.Sprintf("%s/%s/%s imagePullPolicy=%s", pod.Namespace, pod.Name, container.Name, container.ImagePullPolicy))
					}
				}
			}
		}
		// log for debugging output before we ultimately fail
		e2e.Logf("Pods found with invalid container images not present in release payload: %s", strings.Join(invalidPodContainerImages.List(), "\n"))
		e2e.Logf("Pods found with invalid container image pull policy not equal to IfNotPresent: %s", strings.Join(invalidPodContainerImagePullPolicy.List(), "\n"))
		if len(invalidPodContainerImages) > 0 {
			e2e.Failf("Pods found with invalid container images not present in release payload: %s", strings.Join(invalidPodContainerImages.List(), "\n"))
		}
		if len(invalidPodContainerImagePullPolicy) > 0 {
			e2e.Failf("Pods found with invalid container image pull policy not equal to IfNotPresent: %s", strings.Join(invalidPodContainerImagePullPolicy.List(), "\n"))
		}
	})
})
<|endoftext|>


<commit_before>// +build ignore

package main

import (
	"fmt"

	"v.io/v23"
	"v.io/v23/rpc"
	"v.io/v23/security"
	"v.io/x/lib/vlog"
	"v.io/x/ref/lib/signals"
	_ "v.io/x/ref/profiles"

	"pingpong"
)

type pongd struct{}

func (f *pongd) Ping(call rpc.ServerCall, message string) (result string, err error) {
	remote, _ := security.BlessingNames(call.Context(), security.CallSideRemote)
	fmt.Printf("%v: %q\n", remote, message)
	return "PONG", nil
}

func main() {
	ctx, shutdown := v23.Init()
	defer shutdown()

	s, err := v23.NewServer(ctx)
	if err != nil {
		vlog.Fatal("failure creating server: ", err)
	}

	serverPong := pingpong.PingPongServer(&pongd{})

	fmt.Printf("Starting server\n")
	if endpoint, err := s.Listen(v23.GetListenSpec(ctx)); err == nil {
		fmt.Printf("Listening at: %v\n", endpoint)
	} else {
		vlog.Fatal("error listening to service: ", err)
	}

	if err := s.Serve("pingpong", serverPong, nil); err == nil {
		fmt.Printf("Serving pingpong\n")
	} else {
		vlog.Fatal("error serving service: ", err)
	}

	// Wait forever.
	<-signals.ShutdownOnSignals(ctx)
}
<commit_msg>"projects/playground": Use security.[Local/Remote]BlessingNames<commit_after>// +build ignore

package main

import (
	"fmt"

	"v.io/v23"
	"v.io/v23/rpc"
	"v.io/v23/security"
	"v.io/x/lib/vlog"
	"v.io/x/ref/lib/signals"
	_ "v.io/x/ref/profiles"

	"pingpong"
)

type pongd struct{}

func (f *pongd) Ping(call rpc.ServerCall, message string) (result string, err error) {
	remote, _ := security.RemoteBlessingNames(call.Context())
	fmt.Printf("%v: %q\n", remote, message)
	return "PONG", nil
}

func main() {
	ctx, shutdown := v23.Init()
	defer shutdown()

	s, err := v23.NewServer(ctx)
	if err != nil {
		vlog.Fatal("failure creating server: ", err)
	}

	serverPong := pingpong.PingPongServer(&pongd{})

	fmt.Printf("Starting server\n")
	if endpoint, err := s.Listen(v23.GetListenSpec(ctx)); err == nil {
		fmt.Printf("Listening at: %v\n", endpoint)
	} else {
		vlog.Fatal("error listening to service: ", err)
	}

	if err := s.Serve("pingpong", serverPong, nil); err == nil {
		fmt.Printf("Serving pingpong\n")
	} else {
		vlog.Fatal("error serving service: ", err)
	}

	// Wait forever.
	<-signals.ShutdownOnSignals(ctx)
}
<|endoftext|>


<commit_before>package payment

import (
	"errors"
	"fmt"
	"socialapi/workers/payment/paymenterrors"
	"socialapi/workers/payment/paymentmodels"
	"socialapi/workers/payment/stripe"
	"time"
)

var (
	ProviderNotFound       = errors.New("provider not found")
	ProviderNotImplemented = errors.New("provider not implemented")
)

//----------------------------------------------------------
// SubscribeRequest
//----------------------------------------------------------

type SubscribeRequest struct {
	AccountId, Token, Email           string
	Provider, PlanTitle, PlanInterval string
}

func (s *SubscribeRequest) Do() (interface{}, error) {
	switch s.Provider {
	case "stripe":
		return nil, stripe.Subscribe(
			s.Token, s.AccountId, s.Email, s.PlanTitle, s.PlanInterval,
		)
	case "paypal":
		return nil, ProviderNotImplemented
	default:
		return nil, ProviderNotFound
	}
}

//----------------------------------------------------------
// SubscriptionRequest
//----------------------------------------------------------

type SubscriptionRequest struct {
	AccountId string
}

type SubscriptionsResponse struct {
	AccountId          string    `json:"accountId"`
	PlanTitle          string    `json:"planTitle"`
	PlanInterval       string    `json:"planInterval"`
	State              string    `json:"state"`
	CurrentPeriodStart time.Time `json:"currentPeriodStart"`
	CurrentPeriodEnd   time.Time `json:"currentPeriodEnd"`
}

// Do checks if given `account_id` is a paying customer and returns
// the current plan the current is subscribed if any.
//
// Errors:
//		paymenterrors.ErrCustomerNotFound if user is found
//		paymenterrors.ErrCustomerNotSubscribedToAnyPlans if user no subscriptions
//		paymenterrors.ErrPlanNotFound if user subscription's plan isn't found
func (s *SubscriptionRequest) Do() (*SubscriptionsResponse, error) {
	if s.AccountId == "" {
		return nil, paymenterrors.ErrAccountIdIsNotSet
	}

	customer, err := stripe.FindCustomerByOldId(s.AccountId)
	if err != nil {
		return nil, err
	}

	subscriptions, err := stripe.FindCustomerActiveSubscriptions(customer)
	if err != nil {
		return nil, err
	}

	if len(subscriptions) == 0 {
		return nil, paymenterrors.ErrCustomerNotSubscribedToAnyPlans
	}

	currentSubscription := subscriptions[0]

	plan := &paymentmodel.Plan{}
	err = plan.ById(currentSubscription.PlanId)
	if err != nil {
		return nil, err
	}

	resp := &SubscriptionsResponse{
		AccountId:          s.AccountId,
		PlanTitle:          plan.Title,
		PlanInterval:       plan.Interval,
		CurrentPeriodStart: currentSubscription.CurrentPeriodStart,
		CurrentPeriodEnd:   currentSubscription.CurrentPeriodEnd,
		State:              currentSubscription.State,
	}

	return resp, nil
}

// DoWithDefault is different from Do since client excepts to get
// "free" plan regardless of user not found or doesn't have any
// subscriptions etc.
func (s *SubscriptionRequest) DoWithDefault() (*SubscriptionsResponse, error) {
	resp, err := s.Do()
	if err == nil {
		return resp, nil
	}

	defaultResp := &SubscriptionsResponse{
		AccountId:    s.AccountId,
		PlanTitle:    "free",
		PlanInterval: "month",
		State:        "active",
	}

	defaultResponseErrs := []error{
		paymenterrors.ErrCustomerNotSubscribedToAnyPlans,
		paymenterrors.ErrCustomerNotFound,
		paymenterrors.ErrPlanNotFound,
	}

	for _, respError := range defaultResponseErrs {
		if err == respError {
			return defaultResp, nil
		}
	}

	return nil, err
}

//----------------------------------------------------------
// InvoiceRequest
//----------------------------------------------------------

type InvoiceRequest struct {
	AccountId string
}

func (i *InvoiceRequest) Do() ([]*stripe.StripeInvoiceResponse, error) {
	invoices, err := stripe.FindInvoicesForCustomer(i.AccountId)
	if err != nil {
		return nil, err
	}

	return invoices, nil
}

//----------------------------------------------------------
// GetCreditCard
//----------------------------------------------------------

type CreditCardRequest struct {
	AccountId string
}

func (c *CreditCardRequest) Do() (*stripe.CreditCardResponse, error) {
	resp, err := stripe.GetCreditCard(c.AccountId)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

//----------------------------------------------------------
// UpdateCreditCard
//----------------------------------------------------------

type UpdateCreditCardRequest struct {
	AccountId, Provider, Token string
}

func (u *UpdateCreditCardRequest) Do() (interface{}, error) {
	switch u.Provider {
	case "stripe":
		return nil, stripe.UpdateCreditCard(u.AccountId, u.Token)
	case "paypal":
		return nil, ProviderNotImplemented
	default:
		return nil, ProviderNotFound
	}
}

//----------------------------------------------------------
// StripeWebhook
//----------------------------------------------------------

type StripeWebhook struct {
	Name     string `json:"type"`
	Created  int    `json:"created"`
	Livemode bool   `json:"livemode"`
	Id       string `json:"id"`
	Data     []byte `json:"data"`
	Object   string `json:"object"`
}

func (s *StripeWebhook) Do() (interface{}, error) {
	var err error

	switch s.Name {
	case "customer.subscription.deleted":
		err = stripe.SubscriptionDeletedWebhook(s.Data)
	default:
		fmt.Println("Unhandled Stripe webhook", s.Name)
	}

	return nil, err
}
<commit_msg>payment: fix unmarshalling webhook json<commit_after>package payment

import (
	"encoding/json"
	"errors"
	"fmt"
	"socialapi/workers/payment/paymenterrors"
	"socialapi/workers/payment/paymentmodels"
	"socialapi/workers/payment/stripe"
	"time"
)

var (
	ProviderNotFound       = errors.New("provider not found")
	ProviderNotImplemented = errors.New("provider not implemented")
)

//----------------------------------------------------------
// SubscribeRequest
//----------------------------------------------------------

type SubscribeRequest struct {
	AccountId, Token, Email           string
	Provider, PlanTitle, PlanInterval string
}

func (s *SubscribeRequest) Do() (interface{}, error) {
	switch s.Provider {
	case "stripe":
		return nil, stripe.Subscribe(
			s.Token, s.AccountId, s.Email, s.PlanTitle, s.PlanInterval,
		)
	case "paypal":
		return nil, ProviderNotImplemented
	default:
		return nil, ProviderNotFound
	}
}

//----------------------------------------------------------
// SubscriptionRequest
//----------------------------------------------------------

type SubscriptionRequest struct {
	AccountId string
}

type SubscriptionsResponse struct {
	AccountId          string    `json:"accountId"`
	PlanTitle          string    `json:"planTitle"`
	PlanInterval       string    `json:"planInterval"`
	State              string    `json:"state"`
	CurrentPeriodStart time.Time `json:"currentPeriodStart"`
	CurrentPeriodEnd   time.Time `json:"currentPeriodEnd"`
}

// Do checks if given `account_id` is a paying customer and returns
// the current plan the current is subscribed if any.
//
// Errors:
//		paymenterrors.ErrCustomerNotFound if user is found
//		paymenterrors.ErrCustomerNotSubscribedToAnyPlans if user no subscriptions
//		paymenterrors.ErrPlanNotFound if user subscription's plan isn't found
func (s *SubscriptionRequest) Do() (*SubscriptionsResponse, error) {
	if s.AccountId == "" {
		return nil, paymenterrors.ErrAccountIdIsNotSet
	}

	customer, err := stripe.FindCustomerByOldId(s.AccountId)
	if err != nil {
		return nil, err
	}

	subscriptions, err := stripe.FindCustomerActiveSubscriptions(customer)
	if err != nil {
		return nil, err
	}

	if len(subscriptions) == 0 {
		return nil, paymenterrors.ErrCustomerNotSubscribedToAnyPlans
	}

	currentSubscription := subscriptions[0]

	plan := &paymentmodel.Plan{}
	err = plan.ById(currentSubscription.PlanId)
	if err != nil {
		return nil, err
	}

	resp := &SubscriptionsResponse{
		AccountId:          s.AccountId,
		PlanTitle:          plan.Title,
		PlanInterval:       plan.Interval,
		CurrentPeriodStart: currentSubscription.CurrentPeriodStart,
		CurrentPeriodEnd:   currentSubscription.CurrentPeriodEnd,
		State:              currentSubscription.State,
	}

	return resp, nil
}

// DoWithDefault is different from Do since client excepts to get
// "free" plan regardless of user not found or doesn't have any
// subscriptions etc.
func (s *SubscriptionRequest) DoWithDefault() (*SubscriptionsResponse, error) {
	resp, err := s.Do()
	if err == nil {
		return resp, nil
	}

	defaultResp := &SubscriptionsResponse{
		AccountId:    s.AccountId,
		PlanTitle:    "free",
		PlanInterval: "month",
		State:        "active",
	}

	defaultResponseErrs := []error{
		paymenterrors.ErrCustomerNotSubscribedToAnyPlans,
		paymenterrors.ErrCustomerNotFound,
		paymenterrors.ErrPlanNotFound,
	}

	for _, respError := range defaultResponseErrs {
		if err == respError {
			return defaultResp, nil
		}
	}

	return nil, err
}

//----------------------------------------------------------
// InvoiceRequest
//----------------------------------------------------------

type InvoiceRequest struct {
	AccountId string
}

func (i *InvoiceRequest) Do() ([]*stripe.StripeInvoiceResponse, error) {
	invoices, err := stripe.FindInvoicesForCustomer(i.AccountId)
	if err != nil {
		return nil, err
	}

	return invoices, nil
}

//----------------------------------------------------------
// GetCreditCard
//----------------------------------------------------------

type CreditCardRequest struct {
	AccountId string
}

func (c *CreditCardRequest) Do() (*stripe.CreditCardResponse, error) {
	resp, err := stripe.GetCreditCard(c.AccountId)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

//----------------------------------------------------------
// UpdateCreditCard
//----------------------------------------------------------

type UpdateCreditCardRequest struct {
	AccountId, Provider, Token string
}

func (u *UpdateCreditCardRequest) Do() (interface{}, error) {
	switch u.Provider {
	case "stripe":
		return nil, stripe.UpdateCreditCard(u.AccountId, u.Token)
	case "paypal":
		return nil, ProviderNotImplemented
	default:
		return nil, ProviderNotFound
	}
}

//----------------------------------------------------------
// StripeWebhook
//----------------------------------------------------------

type StripeWebhook struct {
	Name     string `json:"type"`
	Created  int    `json:"created"`
	Livemode bool   `json:"livemode"`
	Id       string `json:"id"`
	Data     struct {
		Object interface{} `json:"object"`
	} `json:"data"`
}

func (s *StripeWebhook) Do() (interface{}, error) {
	var err error

	switch s.Name {
	case "customer.subscription.deleted":
		raw, err := json.Marshal(s.Data.Object)
		if err != nil {
			return nil, err
		}

		err = stripe.SubscriptionDeletedWebhook(raw)
	default:
		fmt.Println("Unhandled Stripe webhook", s.Name)
	}

	return nil, err
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
)

func main() {
	fmt.Println("When's Saturday?")
	today := time.Now().Weekday()

	switch time.Saturday {
	case today + 0:
		fmt.Println("Today.")
	case today + 1:
		fmt.Println("Tommorow.")
	case today + 2:
		fmt.Println("In tow days.")
	default:
		fmt.Println("Too far away.")
	}
}
<commit_msg>Updated 46_SwitchEvaluationOrder.<commit_after>package main

import (
	"fmt"
	"time"
)

func main() {
	fmt.Println("When's Sunday?")
	today := time.Now().Weekday()

	switch time.Sunday {
	case today + 0:
		fmt.Println("Today.")
	case today + 1:
		fmt.Println("Tommorow.")
	case today + 2:
		fmt.Println("In tow days.")
	default:
		fmt.Println("Too far away.")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"math/rand"
	"os"
	"time"

	"github.com/codegangsta/cli"
)

func main() {
	app := cli.NewApp()
	app.Name = "go-outkast"
	app.Usage = "test outkast"
	app.Action = func(c *cli.Context) {
		// Grab all the arguments
		ca := c.Args()

		// Get the total number of arguments
		l := len(c.Args())

		// Select a random item from args
		removedItem := random(0, l)

		// Remove random item from ca
		ca = append(ca[:removedItem], ca[removedItem+1:]...)

		for i := range ca {
			print(c.Args()[i], " ")
		}
	}

	app.Run(os.Args)
}

// Random number generator function 2
func random(min, max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max-min) + min
}
<commit_msg>Updating API to be able to remove a specific number of items<commit_after>package main

import (
	"fmt"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/codegangsta/cli"
)

func main() {
	app := cli.NewApp()
	app.Name = "go-outkast"
	app.Usage = "A fun commandline utility that allows you to randomly remove item(s) from a list of args."
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "rand, r",
			Value: "1",
			Usage: "language for the greeting",
		},
	}
	app.Action = func(c *cli.Context) {
		// Grab all the arguments
		ca := c.Args()

		// Variables to use later
		var (
			rc  int
			err error
		)

		if c.String("rand") != "" {
			r := c.String("rand")
			rc, err = strconv.Atoi(r)
			if err != nil {
				// handle error
				fmt.Println(err)
				os.Exit(2)
			}
		}

		// Get the total number of arguments
		l := len(c.Args())

		if rc > 1 && rc < len(c.Args()) {
			for i := 0; i < rc; i++ {
				ca, l = Remove(ca, l)
			}
		} else {
			ca, l = Remove(ca, l)
		}

		for i := range ca {
			print(c.Args()[i], " ")
		}
	}

	app.Run(os.Args)
}

// Remove -
func Remove(ca []string, l int) ([]string, int) {
	// Select a random item from args
	removedItem := random(0, l)

	// Remove random item from ca
	ca = append(ca[:removedItem], ca[removedItem+1:]...)
	// Reset the length of the slice in order to return it
	l = len(ca)
	return ca, l
}

// Random number generator function 2
func random(min, max int) int {
	rand.Seed(time.Now().Unix())
	return rand.Intn(max-min) + min
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"os"
	"path"
	"strconv"

	"github.com/subgraph/oz/oz-daemon"
	"github.com/subgraph/oz/oz-init"

	"github.com/codegangsta/cli"
)

type fnRunType func()

var runFunc fnRunType
var runBasename string

func init() {
	runBasename = path.Base(os.Args[0])

	switch runBasename {
	case "oz":
		runFunc = runApplication
	default:
		// TODO: Exit if already inside sandbox should only happen
		runFunc = runSandbox
	}
}

func main() {
	runFunc()
}

func runSandbox() {
	hostname, _ := os.Hostname()
	if runBasename == hostname {
		fmt.Fprintf(os.Stderr, "Cannot launch from inside a sandbox.\n")
		os.Exit(1)
	}

	err := daemon.Launch(runBasename, os.Args[1:], os.Environ(), false)
	if err != nil {
		fmt.Fprintf(os.Stderr, "launch command failed: %v.\n", err)
		os.Exit(1)
	}
}

func runApplication() {
	app := cli.NewApp()

	app.Name = "oz"
	app.Usage = "command line interface to Oz sandboxes"
	app.Author = "Subgraph"
	app.Email = "info@subgraph.com"
	app.Version = OzVersion
	app.EnableBashCompletion = true
	app.Commands = []cli.Command{
		{
			Name:   "profiles",
			Usage:  "list available application profiles",
			Action: handleProfiles,
		},
		{
			Name:   "launch",
			Usage:  "launch an application profile",
			Action: handleLaunch,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name: "noexec, n",
				},
			},
		},
		{
			Name:   "list",
			Usage:  "list running sandboxes",
			Action: handleList,
		},
		{
			Name:   "shell",
			Usage:  "start a shell in a running sandbox",
			Action: handleShell,
		},
		{
			Name:   "clean",
			Action: handleClean,
		},
		{
			Name:   "kill",
			Action: handleKill,
		},
		{
			Name:   "logs",
			Action: handleLogs,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name: "f",
				},
			},
		},
	}
	app.Run(os.Args)
}

func handleProfiles(c *cli.Context) {
	ps, err := daemon.ListProfiles()
	if err != nil {
		fmt.Printf("Error listing profiles: %v\n", err)
		os.Exit(1)
	}
	for i, p := range ps {
		fmt.Printf("%2d) %-30s %s\n", i+1, p.Name, p.Path)
	}
}

func handleLaunch(c *cli.Context) {
	noexec := c.Bool("noexec")
	if len(c.Args()) == 0 {
		fmt.Println("Argument needed to launch command")
		os.Exit(1)
	}
	err := daemon.Launch(c.Args()[0], c.Args()[1:], os.Environ(), noexec)
	if err != nil {
		fmt.Printf("launch command failed: %v\n", err)
		os.Exit(1)
	}
}

func handleList(c *cli.Context) {
	sboxes, err := daemon.ListSandboxes()
	if err != nil {
		fmt.Printf("Error listing running sandboxes: %v\n", err)
		os.Exit(1)
	}
	if len(sboxes) == 0 {
		fmt.Println("No running sandboxes")
		return
	}
	for _, sb := range sboxes {
		fmt.Printf("%2d) %s\n", sb.Id, sb.Profile)

	}
}

func handleShell(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("Sandbox id argument needed")
		os.Exit(1)
	}
	id, err := strconv.Atoi(c.Args()[0])
	if err != nil {
		fmt.Println("Sandbox id argument must be an integer")
		os.Exit(1)
	}

	sb, err := getSandboxById(id)
	if err != nil {
		fmt.Printf("Error retrieving sandbox list: %v\n", err)
		os.Exit(1)
	}
	if sb == nil {
		fmt.Printf("No sandbox found with id = %d\n", id)
		os.Exit(1)
	}
	term := os.Getenv("TERM")
	fd, err := ozinit.RunShell(sb.Address, term)
	if err != nil {
		fmt.Printf("start shell command failed: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("Entering interactive shell?\n")
	st, err := SetRawTerminal(0)
	HandleResize(fd)
	f := os.NewFile(uintptr(fd), "")
	go io.Copy(f, os.Stdin)
	io.Copy(os.Stdout, f)
	RestoreTerminal(0, st)
	fmt.Println("done..")
}

func getSandboxById(id int) (*daemon.SandboxInfo, error) {
	sboxes, err := daemon.ListSandboxes()
	if err != nil {
		return nil, err
	}
	for _, sb := range sboxes {
		if id == sb.Id {
			return &sb, nil
		}
	}
	return nil, nil
}

func handleClean(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("Need a profile to clean")
		os.Exit(1)
	}
	err := daemon.Clean(c.Args()[0])
	if err != nil {
		fmt.Println("Clean failed:", err)
	}
}

func handleKill(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("Need a sandbox id to kill\n")
		os.Exit(1)
	}
	id, err := strconv.Atoi(c.Args()[0])
	if err != nil {
		fmt.Printf("Could not parse id value %s\n", c.Args()[0])
		os.Exit(1)
	}
	if err := daemon.KillSandbox(id); err != nil {
		fmt.Println("Kill command failed:", err)
	}

}
func handleLogs(c *cli.Context) {
	follow := c.Bool("f")
	ch, err := daemon.Logs(0, follow)
	if err != nil {
		fmt.Println("Logs failed", err)
		os.Exit(1)
	}
	for ll := range ch {
		fmt.Println(ll)
	}
}
<commit_msg>Wooops<commit_after>package main

import (
	"fmt"
	"io"
	"os"
	"path"
	"strconv"

	"github.com/subgraph/oz"
	"github.com/subgraph/oz/oz-daemon"
	"github.com/subgraph/oz/oz-init"

	"github.com/codegangsta/cli"
)

type fnRunType func()

var runFunc fnRunType
var runBasename string

func init() {
	runBasename = path.Base(os.Args[0])

	switch runBasename {
	case "oz":
		runFunc = runApplication
	default:
		// TODO: Exit if already inside sandbox should only happen
		runFunc = runSandbox
	}
}

func main() {
	runFunc()
}

func runSandbox() {
	hostname, _ := os.Hostname()
	if runBasename == hostname {
		fmt.Fprintf(os.Stderr, "Cannot launch from inside a sandbox.\n")
		os.Exit(1)
	}

	err := daemon.Launch(runBasename, os.Args[1:], os.Environ(), false)
	if err != nil {
		fmt.Fprintf(os.Stderr, "launch command failed: %v.\n", err)
		os.Exit(1)
	}
}

func runApplication() {
	app := cli.NewApp()

	app.Name = "oz"
	app.Usage = "command line interface to Oz sandboxes"
	app.Author = "Subgraph"
	app.Email = "info@subgraph.com"
	app.Version = oz.OzVersion
	app.EnableBashCompletion = true
	app.Commands = []cli.Command{
		{
			Name:   "profiles",
			Usage:  "list available application profiles",
			Action: handleProfiles,
		},
		{
			Name:   "launch",
			Usage:  "launch an application profile",
			Action: handleLaunch,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name: "noexec, n",
				},
			},
		},
		{
			Name:   "list",
			Usage:  "list running sandboxes",
			Action: handleList,
		},
		{
			Name:   "shell",
			Usage:  "start a shell in a running sandbox",
			Action: handleShell,
		},
		{
			Name:   "clean",
			Action: handleClean,
		},
		{
			Name:   "kill",
			Action: handleKill,
		},
		{
			Name:   "logs",
			Action: handleLogs,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name: "f",
				},
			},
		},
	}
	app.Run(os.Args)
}

func handleProfiles(c *cli.Context) {
	ps, err := daemon.ListProfiles()
	if err != nil {
		fmt.Printf("Error listing profiles: %v\n", err)
		os.Exit(1)
	}
	for i, p := range ps {
		fmt.Printf("%2d) %-30s %s\n", i+1, p.Name, p.Path)
	}
}

func handleLaunch(c *cli.Context) {
	noexec := c.Bool("noexec")
	if len(c.Args()) == 0 {
		fmt.Println("Argument needed to launch command")
		os.Exit(1)
	}
	err := daemon.Launch(c.Args()[0], c.Args()[1:], os.Environ(), noexec)
	if err != nil {
		fmt.Printf("launch command failed: %v\n", err)
		os.Exit(1)
	}
}

func handleList(c *cli.Context) {
	sboxes, err := daemon.ListSandboxes()
	if err != nil {
		fmt.Printf("Error listing running sandboxes: %v\n", err)
		os.Exit(1)
	}
	if len(sboxes) == 0 {
		fmt.Println("No running sandboxes")
		return
	}
	for _, sb := range sboxes {
		fmt.Printf("%2d) %s\n", sb.Id, sb.Profile)

	}
}

func handleShell(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("Sandbox id argument needed")
		os.Exit(1)
	}
	id, err := strconv.Atoi(c.Args()[0])
	if err != nil {
		fmt.Println("Sandbox id argument must be an integer")
		os.Exit(1)
	}

	sb, err := getSandboxById(id)
	if err != nil {
		fmt.Printf("Error retrieving sandbox list: %v\n", err)
		os.Exit(1)
	}
	if sb == nil {
		fmt.Printf("No sandbox found with id = %d\n", id)
		os.Exit(1)
	}
	term := os.Getenv("TERM")
	fd, err := ozinit.RunShell(sb.Address, term)
	if err != nil {
		fmt.Printf("start shell command failed: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("Entering interactive shell?\n")
	st, err := SetRawTerminal(0)
	HandleResize(fd)
	f := os.NewFile(uintptr(fd), "")
	go io.Copy(f, os.Stdin)
	io.Copy(os.Stdout, f)
	RestoreTerminal(0, st)
	fmt.Println("done..")
}

func getSandboxById(id int) (*daemon.SandboxInfo, error) {
	sboxes, err := daemon.ListSandboxes()
	if err != nil {
		return nil, err
	}
	for _, sb := range sboxes {
		if id == sb.Id {
			return &sb, nil
		}
	}
	return nil, nil
}

func handleClean(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("Need a profile to clean")
		os.Exit(1)
	}
	err := daemon.Clean(c.Args()[0])
	if err != nil {
		fmt.Println("Clean failed:", err)
	}
}

func handleKill(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("Need a sandbox id to kill\n")
		os.Exit(1)
	}
	id, err := strconv.Atoi(c.Args()[0])
	if err != nil {
		fmt.Printf("Could not parse id value %s\n", c.Args()[0])
		os.Exit(1)
	}
	if err := daemon.KillSandbox(id); err != nil {
		fmt.Println("Kill command failed:", err)
	}

}
func handleLogs(c *cli.Context) {
	follow := c.Bool("f")
	ch, err := daemon.Logs(0, follow)
	if err != nil {
		fmt.Println("Logs failed", err)
		os.Exit(1)
	}
	for ll := range ch {
		fmt.Println(ll)
	}
}
<|endoftext|>


<commit_before>package authentication

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/Azure/go-autorest/autorest/azure"
)

type Environment struct {
	Portal                  string         `json:"portal"`
	Authentication          Authentication `json:"authentication"`
	Media                   string         `json:"media"`
	GraphAudience           string         `json:"graphAudience"`
	Graph                   string         `json:"graph"`
	Name                    string         `json:"name"`
	Suffixes                Suffixes       `json:"suffixes"`
	Batch                   string         `json:"batch"`
	ResourceManager         string         `json:"resourceManager"`
	VmImageAliasDoc         string         `json:"vmImageAliasDoc"`
	ActiveDirectoryDataLake string         `json:"activeDirectoryDataLake"`
	SqlManagement           string         `json:"sqlManagement"`
	Gallery                 string         `json:"gallery"`
}

type Authentication struct {
	LoginEndpoint    string   `json:"loginEndpoint"`
	Audiences        []string `json:"audiences"`
	Tenant           string   `json:"tenant"`
	IdentityProvider string   `json:"identityProvider"`
}

type Suffixes struct {
	AzureDataLakeStoreFileSystem        string `json:"azureDataLakeStoreFileSystem"`
	AcrLoginServer                      string `json:"acrLoginServer"`
	SqlServerHostname                   string `json:"sqlServerHostname"`
	AzureDataLakeAnalyticsCatalogAndJob string `json:"azureDataLakeAnalyticsCatalogAndJob"`
	KeyVaultDns                         string `json:"keyVaultDns"`
	Storage                             string `json:"storage"`
	AzureFrontDoorEndpointSuffix        string `json:"azureFrontDoorEndpointSuffix"`
}

// DetermineEnvironment determines what the Environment name is within
// the Azure SDK for Go and then returns the association environment, if it exists.
func DetermineEnvironment(name string) (*azure.Environment, error) {
	// detect cloud from environment
	env, envErr := azure.EnvironmentFromName(name)

	if envErr != nil {
		// try again with wrapped value to support readable values like german instead of AZUREGERMANCLOUD
		wrapped := fmt.Sprintf("AZURE%sCLOUD", name)
		env, envErr = azure.EnvironmentFromName(wrapped)
		if envErr != nil {
			return nil, fmt.Errorf("An Azure Environment with name %q was not found: %+v", name, envErr)
		}
	}

	return &env, nil
}

// LoadEnvironmentFromUrl attempts to load the specified environment from the endpoint.
// if the endpoint is an empty string, or an environment can't be
// found at the endpoint url then an error is returned
func LoadEnvironmentFromUrl(endpoint string) (*azure.Environment, error) {
	if endpoint == "" {
		return nil, fmt.Errorf("Endpoint was not set!")
	}

	env, err := azure.EnvironmentFromURL(endpoint)
	if err != nil {
		return nil, fmt.Errorf("Error retrieving Environment from Endpoint %q: %+v", endpoint, err)
	}

	return &env, nil
}

func normalizeEnvironmentName(input string) string {
	// Environment is stored as `Azure{Environment}Cloud`
	output := strings.ToLower(input)
	output = strings.TrimPrefix(output, "azure")
	output = strings.TrimSuffix(output, "cloud")

	// however Azure Public is `AzureCloud` in the CLI Profile and not `AzurePublicCloud`.
	if output == "" {
		return "public"
	}
	return output
}

// AzureEnvironmentByName returns a specific Azure Environment from the specified endpoint
func AzureEnvironmentByNameFromEndpoint(ctx context.Context, endpoint string, environmentName string) (*azure.Environment, error) {
	var environmentTranslationMap = map[string]string{
		"public": "AzureCloud",
		"usgovernment": "AzureUSGovernment",
		"german": "AzureGermanCloud",
		"china": "AzureChinaCloud",
	}

	uri := fmt.Sprintf("https://%s//metadata/endpoints?api-version=2019-05-01", endpoint)
	client := http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
		},
	}
	req, err := http.NewRequestWithContext(ctx, "GET", uri, nil)
	if err != nil {
		return nil, err
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("retrieving environments from Azure MetaData service: %+v", err)
	}

	dec := json.NewDecoder(resp.Body)
	if _, err = dec.Token(); err != nil {
		return nil, err
	}

	// while the array contains values
	for dec.More() {
		var env Environment
		// decode an array value (Message)
		err := dec.Decode(&env)
		if err != nil {
			return nil, fmt.Errorf("unable to decode environment from %q response: %+v", uri, err)
		}
		if strings.EqualFold(env.Name, environmentName) || strings.EqualFold(env.Name, environmentTranslationMap[environmentName]) {
			// This check is to prevent azure stack users from getting their environment this way
			if env.Authentication.Tenant != "common" && env.Authentication.IdentityProvider != "AAD" {
				return nil, fmt.Errorf("environment %q from metadata_url %q is not supported", environmentName, endpoint)
			}

			aEnv := &azure.Environment{
				ResourceManagerEndpoint: env.ResourceManager,
				StorageEndpointSuffix:   env.Suffixes.Storage,
				ActiveDirectoryEndpoint: env.Authentication.LoginEndpoint,
				GraphEndpoint:           env.Graph,
				ResourceIdentifiers: azure.ResourceIdentifier{
					// This isn't returned from the metadata url and is universal across all environments
					Storage: "https://storage.azure.com/",
				},
			}

			if len(env.Authentication.Audiences) > 0 {
				aEnv.TokenAudience = env.Authentication.Audiences[0]
			} else {
				return nil, fmt.Errorf("unable to find token audience for environment %q from endpoint %q", environmentName, endpoint)
			}

			return aEnv, nil
		}
	}

	return nil, fmt.Errorf("unable to find environment %q from endpoint %q", environmentName, endpoint)
}
<commit_msg>Address review<commit_after>package authentication

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/Azure/go-autorest/autorest/azure"
)

type Environment struct {
	Portal                  string         `json:"portal"`
	Authentication          Authentication `json:"authentication"`
	Media                   string         `json:"media"`
	GraphAudience           string         `json:"graphAudience"`
	Graph                   string         `json:"graph"`
	Name                    string         `json:"name"`
	Suffixes                Suffixes       `json:"suffixes"`
	Batch                   string         `json:"batch"`
	ResourceManager         string         `json:"resourceManager"`
	VmImageAliasDoc         string         `json:"vmImageAliasDoc"`
	ActiveDirectoryDataLake string         `json:"activeDirectoryDataLake"`
	SqlManagement           string         `json:"sqlManagement"`
	Gallery                 string         `json:"gallery"`
}

type Authentication struct {
	LoginEndpoint    string   `json:"loginEndpoint"`
	Audiences        []string `json:"audiences"`
	Tenant           string   `json:"tenant"`
	IdentityProvider string   `json:"identityProvider"`
}

type Suffixes struct {
	AzureDataLakeStoreFileSystem        string `json:"azureDataLakeStoreFileSystem"`
	AcrLoginServer                      string `json:"acrLoginServer"`
	SqlServerHostname                   string `json:"sqlServerHostname"`
	AzureDataLakeAnalyticsCatalogAndJob string `json:"azureDataLakeAnalyticsCatalogAndJob"`
	KeyVaultDns                         string `json:"keyVaultDns"`
	Storage                             string `json:"storage"`
	AzureFrontDoorEndpointSuffix        string `json:"azureFrontDoorEndpointSuffix"`
}

// DetermineEnvironment determines what the Environment name is within
// the Azure SDK for Go and then returns the association environment, if it exists.
func DetermineEnvironment(name string) (*azure.Environment, error) {
	// detect cloud from environment
	env, envErr := azure.EnvironmentFromName(name)

	if envErr != nil {
		// try again with wrapped value to support readable values like german instead of AZUREGERMANCLOUD
		wrapped := fmt.Sprintf("AZURE%sCLOUD", name)
		env, envErr = azure.EnvironmentFromName(wrapped)
		if envErr != nil {
			return nil, fmt.Errorf("An Azure Environment with name %q was not found: %+v", name, envErr)
		}
	}

	return &env, nil
}

// LoadEnvironmentFromUrl attempts to load the specified environment from the endpoint.
// if the endpoint is an empty string, or an environment can't be
// found at the endpoint url then an error is returned
func LoadEnvironmentFromUrl(endpoint string) (*azure.Environment, error) {
	if endpoint == "" {
		return nil, fmt.Errorf("Endpoint was not set!")
	}

	env, err := azure.EnvironmentFromURL(endpoint)
	if err != nil {
		return nil, fmt.Errorf("Error retrieving Environment from Endpoint %q: %+v", endpoint, err)
	}

	return &env, nil
}

func normalizeEnvironmentName(input string) string {
	// Environment is stored as `Azure{Environment}Cloud`
	output := strings.ToLower(input)
	output = strings.TrimPrefix(output, "azure")
	output = strings.TrimSuffix(output, "cloud")

	// however Azure Public is `AzureCloud` in the CLI Profile and not `AzurePublicCloud`.
	if output == "" {
		return "public"
	}
	return output
}

// AzureEnvironmentByName returns a specific Azure Environment from the specified endpoint
func AzureEnvironmentByNameFromEndpoint(ctx context.Context, endpoint string, environmentName string) (*azure.Environment, error) {
	var environmentTranslationMap = map[string]string{
		"public": "AzureCloud",
		"usgovernment": "AzureUSGovernment",
		"german": "AzureGermanCloud",
		"china": "AzureChinaCloud",
	}

	uri := fmt.Sprintf("https://%s/metadata/endpoints?api-version=2019-05-01", endpoint)
	client := http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
		},
	}
	req, err := http.NewRequestWithContext(ctx, "GET", uri, nil)
	if err != nil {
		return nil, err
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("retrieving environments from Azure MetaData service: %+v", err)
	}

	dec := json.NewDecoder(resp.Body)
	if _, err = dec.Token(); err != nil {
		return nil, err
	}

	// while the array contains values
	for dec.More() {
		var env Environment
		// decode an array value (Message)
		err := dec.Decode(&env)
		if err != nil {
			return nil, fmt.Errorf("unable to decode environment from %q response: %+v", uri, err)
		}
		if strings.EqualFold(env.Name, environmentName) || strings.EqualFold(env.Name, environmentTranslationMap[environmentName]) {
			// This check is to prevent azure stack users from getting their environment this way
			if env.Authentication.Tenant != "common" && env.Authentication.IdentityProvider != "AAD" {
				return nil, fmt.Errorf("environment %q from metadata_url %q is not supported", environmentName, endpoint)
			}

			aEnv := &azure.Environment{
				ResourceManagerEndpoint: env.ResourceManager,
				StorageEndpointSuffix:   env.Suffixes.Storage,
				ActiveDirectoryEndpoint: env.Authentication.LoginEndpoint,
				GraphEndpoint:           env.Graph,
				ResourceIdentifiers: azure.ResourceIdentifier{
					// This isn't returned from the metadata url and is universal across all environments
					Storage: "https://storage.azure.com/",
				},
			}

			if len(env.Authentication.Audiences) > 0 {
				aEnv.TokenAudience = env.Authentication.Audiences[0]
			} else {
				return nil, fmt.Errorf("unable to find token audience for environment %q from endpoint %q", environmentName, endpoint)
			}

			return aEnv, nil
		}
	}

	return nil, fmt.Errorf("unable to find environment %q from endpoint %q", environmentName, endpoint)
}
<|endoftext|>


<commit_before>// Clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	target := time.Date(2018, 9, 22, 21, 54, 0, 0, time.Local)
	motto := "Just Go"
	printTargetTime(target, motto)
	exitOnEnterKey()

	var previous time.Time
	for {
		now := time.Now().Truncate(time.Second)
		if now != previous {
			previous = now
			countdown := now.Sub(target) // Negative times are before the target
			printCountdown(now, countdown)
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func exitOnEnterKey() {
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()
}

const (
	highlightStart = "\x1b[1;35m"
	highlightEnd   = "\x1b[0m"
	indent         = "\t"
)

func printTargetTime(target time.Time, motto string) {
	fmt.Print(indent, highlightStart, motto, highlightEnd, "\n")
	fmt.Print(indent, target.Format(time.UnixDate), "\n")
}

func printCountdown(now time.Time, countdown time.Duration) {
	var sign string
	if countdown >= 0 {
		sign = "+"
	} else {
		sign = "-"
		countdown = -countdown
	}

	days := int(countdown / (24 * time.Hour))
	countdown = countdown % (24 * time.Hour)

	fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
	if days > 0 {
		fmt.Print(days, "d")
	}
	fmt.Print(countdown, "          \r")
	os.Stdout.Sync()
}
<commit_msg>Equinox day, not minute<commit_after>// Clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	target := time.Date(2018, 9, 22, 0, 0, 0, 0, time.Local)
	motto := "Just Go"
	printTargetTime(target, motto)
	exitOnEnterKey()

	var previous time.Time
	for {
		now := time.Now().Truncate(time.Second)
		if now != previous {
			previous = now
			countdown := now.Sub(target) // Negative times are before the target
			printCountdown(now, countdown)
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func exitOnEnterKey() {
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()
}

const (
	highlightStart = "\x1b[1;35m"
	highlightEnd   = "\x1b[0m"
	indent         = "\t"
)

func printTargetTime(target time.Time, motto string) {
	fmt.Print(indent, highlightStart, motto, highlightEnd, "\n")
	fmt.Print(indent, target.Format(time.UnixDate), "\n")
}

func printCountdown(now time.Time, countdown time.Duration) {
	var sign string
	if countdown >= 0 {
		sign = "+"
	} else {
		sign = "-"
		countdown = -countdown
	}

	days := int(countdown / (24 * time.Hour))
	countdown = countdown % (24 * time.Hour)

	fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
	if days > 0 {
		fmt.Print(days, "d")
	}
	fmt.Print(countdown, "          \r")
	os.Stdout.Sync()
}
<|endoftext|>


<commit_before>/*
 *
 * Copyright 2014 gRPC authors.
 *
 * 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 grpc

import (
	"math"
	"sync"

	"github.com/golang/protobuf/proto"
)

// Codec defines the interface gRPC uses to encode and decode messages.
// Note that implementations of this interface must be thread safe;
// a Codec's methods can be called from concurrent goroutines.
type Codec interface {
	// Marshal returns the wire format of v.
	Marshal(v interface{}) ([]byte, error)
	// Unmarshal parses the wire format into v.
	Unmarshal(data []byte, v interface{}) error
	// String returns the name of the Codec implementation. The returned
	// string will be used as part of content type in transmission.
	String() string
}

// protoCodec is a Codec implementation with protobuf. It is the default codec for gRPC.
type protoCodec struct {
}

type cachedProtoBuffer struct {
	lastMarshaledSize uint32
	proto.Buffer
}

func capToMaxInt32(val int) uint32 {
	if val > math.MaxInt32 {
		return uint32(math.MaxInt32)
	}
	return uint32(val)
}

func (p protoCodec) marshal(v interface{}, cb *cachedProtoBuffer) ([]byte, error) {
	protoMsg := v.(proto.Message)
	newSlice := make([]byte, 0, cb.lastMarshaledSize)

	cb.SetBuf(newSlice)
	cb.Reset()
	if err := cb.Marshal(protoMsg); err != nil {
		return nil, err
	}
	out := cb.Bytes()
	cb.lastMarshaledSize = capToMaxInt32(len(out))
	return out, nil
}

func (p protoCodec) Marshal(v interface{}) ([]byte, error) {
	cb := protoBufferPool.Get().(*cachedProtoBuffer)
	out, err := p.marshal(v, cb)

	// put back buffer and lose the ref to the slice
	cb.SetBuf(nil)
	protoBufferPool.Put(cb)
	return out, err
}

func (p protoCodec) Unmarshal(data []byte, v interface{}) error {
	cb := protoBufferPool.Get().(*cachedProtoBuffer)
	cb.SetBuf(data)
	v.(proto.Message).Reset()
	err := cb.Unmarshal(v.(proto.Message))
	cb.SetBuf(nil)
	protoBufferPool.Put(cb)
	return err
}

func (protoCodec) String() string {
	return "proto"
}

var protoBufferPool = &sync.Pool{
	New: func() interface{} {
		return &cachedProtoBuffer{
			Buffer:            proto.Buffer{},
			lastMarshaledSize: 16,
		}
	},
}
<commit_msg>protoCodec: avoid buffer allocations if proto.Marshaler/Unmarshaler (#1689)<commit_after>/*
 *
 * Copyright 2014 gRPC authors.
 *
 * 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 grpc

import (
	"math"
	"sync"

	"github.com/golang/protobuf/proto"
)

// Codec defines the interface gRPC uses to encode and decode messages.
// Note that implementations of this interface must be thread safe;
// a Codec's methods can be called from concurrent goroutines.
type Codec interface {
	// Marshal returns the wire format of v.
	Marshal(v interface{}) ([]byte, error)
	// Unmarshal parses the wire format into v.
	Unmarshal(data []byte, v interface{}) error
	// String returns the name of the Codec implementation. The returned
	// string will be used as part of content type in transmission.
	String() string
}

// protoCodec is a Codec implementation with protobuf. It is the default codec for gRPC.
type protoCodec struct {
}

type cachedProtoBuffer struct {
	lastMarshaledSize uint32
	proto.Buffer
}

func capToMaxInt32(val int) uint32 {
	if val > math.MaxInt32 {
		return uint32(math.MaxInt32)
	}
	return uint32(val)
}

func (p protoCodec) marshal(v interface{}, cb *cachedProtoBuffer) ([]byte, error) {
	protoMsg := v.(proto.Message)
	newSlice := make([]byte, 0, cb.lastMarshaledSize)

	cb.SetBuf(newSlice)
	cb.Reset()
	if err := cb.Marshal(protoMsg); err != nil {
		return nil, err
	}
	out := cb.Bytes()
	cb.lastMarshaledSize = capToMaxInt32(len(out))
	return out, nil
}

func (p protoCodec) Marshal(v interface{}) ([]byte, error) {
	if pm, ok := v.(proto.Marshaler); ok {
		// object can marshal itself, no need for buffer
		return pm.Marshal()
	}

	cb := protoBufferPool.Get().(*cachedProtoBuffer)
	out, err := p.marshal(v, cb)

	// put back buffer and lose the ref to the slice
	cb.SetBuf(nil)
	protoBufferPool.Put(cb)
	return out, err
}

func (p protoCodec) Unmarshal(data []byte, v interface{}) error {
	protoMsg := v.(proto.Message)
	protoMsg.Reset()

	if pu, ok := protoMsg.(proto.Unmarshaler); ok {
		// object can unmarshal itself, no need for buffer
		return pu.Unmarshal(data)
	}

	cb := protoBufferPool.Get().(*cachedProtoBuffer)
	cb.SetBuf(data)
	err := cb.Unmarshal(protoMsg)
	cb.SetBuf(nil)
	protoBufferPool.Put(cb)
	return err
}

func (protoCodec) String() string {
	return "proto"
}

var protoBufferPool = &sync.Pool{
	New: func() interface{} {
		return &cachedProtoBuffer{
			Buffer:            proto.Buffer{},
			lastMarshaledSize: 16,
		}
	},
}
<|endoftext|>


<commit_before>package boardgame

import (
	"time"
)

type Timer struct {
	//Id will be an opaque identifier that is used to keep track of the
	//corresponding underlying Timer object in the game engine. It is not
	//meaningful to inspect yourself and should not be modified.
	Id       int
	statePtr *state
}

func NewTimer() *Timer {
	return &Timer{}
}

func (t *Timer) Copy() *Timer {
	var result Timer
	result = *t
	return &result
}

func (t *Timer) MarshalJSON() ([]byte, error) {
	obj := map[string]interface{}{
		"Id":       t.Id,
		"TimeLeft": t.TimeLeft(),
	}

	return DefaultMarshalJSON(obj)
}

//Active returns true if the timer is active and counting down.
func (t *Timer) Active() bool {
	return t.Id == 0
}

//TimeLeft returns the number of nanoseconds left until this timer fires.
func (t *Timer) TimeLeft() int {
	if !t.Active() {
		return 0
	}

	return t.statePtr.game.manager.getTimer(t.Id)
}

//Start starts the timer. After nanoseconds nanoseconds have passed, the Move
//will be proposed via proposeMove.
func (t *Timer) Start(nanoseconds int, move Move) {
	game := t.statePtr.game
	manager := game.manager

	t.Id = manager.registerTimer(nanoseconds, game, move)
}

//Cancel cancels an active timer. If the timer is not active, it has no
//effect. Returns true if the timer was active and canceled, false if the
//timer was not active.
func (t *Timer) Cancel() bool {
	if !t.Active() {
		return false
	}
	manager := t.statePtr.game.manager

	manager.cancelTimer(t.Id)

	t.Id = 0

	return true
}

/**************************************
 *
 * GameManager timer-related methods
 *
 **************************************/

//registerTimer registers an implementation level-timer, so that after
//nanoseconds have passed the given move will be proposed to the given game.
//Returns an Id that can be used ot cancel or fetch information on the timer
//in the future. Only designed to be called from Timer.Start()
func (g *GameManager) registerTimer(nanoseconds int, game *Game, move Move) int {
	//TOOD: actually implement this
	return 0
}

//cancelTimer cancels the timer with the given Id. IF there there is no such
//timer, or that timer is already over, it is a no-op.
func (g *GameManager) cancelTimer(id int) {
	//TODO: actually implement this.
}

//getTimer returns the number of nanoseconds left for the timer with the given
//id.
func (g *GameManager) getTimer(id int) int {
	//TODO: actually implement this.
	return 0
}

type timerRecord struct {
	id       int
	index    int
	fireTime time.Time
	game     *Game
	move     Move
}

type timerQueue struct {
	records     []*timerRecord
	recordsById map[int]*timerRecord
}

func (t timerQueue) Len() int {
	return len(t.records)
}

func (t timerQueue) Less(i, j int) bool {
	return t.records[i].fireTime.Sub(t.records[j].fireTime) < 0
}

func (t timerQueue) Swap(i, j int) {
	t.records[i], t.records[j] = t.records[j], t.records[i]
	t.records[i].index = i
	t.records[j].index = j
}

//DO NOT USE THIS DIRECTLY. Use heap.Push(t)
func (t *timerQueue) Push(x interface{}) {
	n := len(t.records)
	item := x.(*timerRecord)
	item.index = n
	t.recordsById[item.id] = item
	t.records = append(t.records, item)
}

//DO NOT USE THIS DIRECTLY. Use heap.Pop()
func (t *timerQueue) Pop() interface{} {
	old := t.records
	n := len(old)
	item := old[n-1]
	item.index = -1
	delete(t.recordsById, item.id)
	t.records = old[0 : n-1]
	return item
}
<commit_msg>A bunch of work to implement timerManager and heaps. Half implemented, not yet tested. Part of #178.<commit_after>package boardgame

import (
	"container/heap"
	"time"
)

type Timer struct {
	//Id will be an opaque identifier that is used to keep track of the
	//corresponding underlying Timer object in the game engine. It is not
	//meaningful to inspect yourself and should not be modified.
	Id       int
	statePtr *state
}

func NewTimer() *Timer {
	return &Timer{}
}

func (t *Timer) Copy() *Timer {
	var result Timer
	result = *t
	return &result
}

func (t *Timer) MarshalJSON() ([]byte, error) {
	obj := map[string]interface{}{
		"Id":       t.Id,
		"TimeLeft": t.TimeLeft(),
	}

	return DefaultMarshalJSON(obj)
}

//Active returns true if the timer is active and counting down.
func (t *Timer) Active() bool {
	return t.Id == 0
}

//TimeLeft returns the number of nanoseconds left until this timer fires.
func (t *Timer) TimeLeft() int {
	if !t.Active() {
		return 0
	}

	return t.statePtr.game.manager.getTimer(t.Id)
}

//Start starts the timer. After nanoseconds nanoseconds have passed, the Move
//will be proposed via proposeMove.
func (t *Timer) Start(nanoseconds int, move Move) {
	game := t.statePtr.game
	manager := game.manager

	t.Id = manager.registerTimer(nanoseconds, game, move)
}

//Cancel cancels an active timer. If the timer is not active, it has no
//effect. Returns true if the timer was active and canceled, false if the
//timer was not active.
func (t *Timer) Cancel() bool {
	if !t.Active() {
		return false
	}
	manager := t.statePtr.game.manager

	manager.cancelTimer(t.Id)

	t.Id = 0

	return true
}

/**************************************
 *
 * GameManager timer-related methods
 *
 **************************************/

//registerTimer registers an implementation level-timer, so that after
//nanoseconds have passed the given move will be proposed to the given game.
//Returns an Id that can be used ot cancel or fetch information on the timer
//in the future. Only designed to be called from Timer.Start()
func (g *GameManager) registerTimer(nanoseconds int, game *Game, move Move) int {
	//TOOD: actually implement this
	return 0
}

//cancelTimer cancels the timer with the given Id. IF there there is no such
//timer, or that timer is already over, it is a no-op.
func (g *GameManager) cancelTimer(id int) {
	//TODO: actually implement this.
}

//getTimer returns the number of nanoseconds left for the timer with the given
//id.
func (g *GameManager) getTimer(id int) int {
	//TODO: actually implement this.
	return 0
}

type timerRecord struct {
	id       int
	index    int
	fireTime time.Time
	game     *Game
	move     Move
}

type timerQueue []*timerRecord

type timerManager struct {
	nextId      int
	records     timerQueue
	recordsById map[int]*timerRecord
}

func newTimerManager() *timerManager {
	return &timerManager{
		nextId:      0,
		records:     make(timerQueue, 0),
		recordsById: make(map[int]*timerRecord),
	}
}

func (t *timerManager) RegisterTimer(nanoseconds int, game *Game, move Move) int {
	record := &timerRecord{
		id:       t.nextId,
		index:    -1,
		fireTime: time.Now(),
		game:     game,
		move:     move,
	}
	t.nextId++

	t.recordsById[record.id] = record

	heap.Push(&t.records, record)

	return record.id
}

//Should be called regularly by the manager to tell this to check and see if
//any timers have fired, and execute them if so.
func (t *timerManager) Tick() {
	for t.nextTimerFired() {
		record := t.popNext()
		if record == nil {
			continue
		}
		if err := <-record.game.ProposeMove(record.move); err != nil {
			//TODO: log the error or something
			panic(err)
		}
	}
}

//Whether the next timer in the queue is already fired
func (t *timerManager) nextTimerFired() bool {
	if len(t.records) == 0 {
		return false
	}

	record := t.records[0]

	if record.fireTime.Sub(time.Now()) < 0 {
		return true
	}

	return false
}

func (t *timerManager) popNext() *timerRecord {
	if !t.nextTimerFired() {
		return nil
	}

	x := heap.Pop(&t.records)

	return x.(*timerRecord)
}

func (t timerQueue) Len() int {
	return len(t)
}

func (t timerQueue) Less(i, j int) bool {
	return t[i].fireTime.Sub(t[j].fireTime) < 0
}

func (t timerQueue) Swap(i, j int) {
	t[i], t[j] = t[j], t[i]
	t[i].index = i
	t[j].index = j
}

//DO NOT USE THIS DIRECTLY. Use heap.Push(t)
func (t *timerQueue) Push(x interface{}) {
	n := len(*t)
	item := x.(*timerRecord)
	item.index = n
	*t = append(*t, item)
}

//DO NOT USE THIS DIRECTLY. Use heap.Pop()
func (t *timerQueue) Pop() interface{} {
	old := *t
	n := len(old)
	item := old[n-1]
	item.index = -1
	*t = old[0 : n-1]
	return item
}
<|endoftext|>


<commit_before>package times

import (
	"errors"
	"time"
)

type Times struct {
	time     time.Time
	timezone string
	location *time.Location
}

// This will return an instance of Times the current time localized to UTC
func Now() *Times {
	timezone := "UTC"
	location, _ := time.LoadLocation(timezone)
	time := time.Now().UTC()
	return &Times{time, timezone, location}
}

// This will localixe
func (t *Times) Normalize(timezone string) (*Times, error) {
	location, err := time.LoadLocation(timezone)

	if err != nil {
		return nil, errors.New("The provided timezone isn't a valid time.Location " + timezone)
	}
	time := t.time.In(location)
	return &Times{time, timezone, location}, nil
}

func (t *Times) Equal(u *Times) bool {
	return t.time.Equal(u.time)
}
<commit_msg>making it so it prints proper<commit_after>package times

import (
	"errors"
	"time"
)

type Times struct {
	time     time.Time
	timezone string
	location *time.Location
}

func (t *Times) String() string {
	return t.time.String()
}

// This will return an instance of Times the current time localized to UTC
func Now() *Times {
	timezone := "UTC"
	location, _ := time.LoadLocation(timezone)
	time := time.Now().UTC()
	return &Times{time, timezone, location}
}

// This will localixe
func (t *Times) Normalize(timezone string) (*Times, error) {
	location, err := time.LoadLocation(timezone)

	if err != nil {
		return nil, errors.New("The provided timezone isn't a valid time.Location " + timezone)
	}
	time := t.time.In(location)
	return &Times{time, timezone, location}, nil
}

func (t *Times) Equal(u *Times) bool {
	return t.time.Equal(u.time)
}
<|endoftext|>


<commit_before>package main

import (
	"zvr/utils"
	"time"
	log "github.com/Sirupsen/logrus"
	"io/ioutil"
	"encoding/json"
	"github.com/pkg/errors"
	"fmt"
	"zvr/server"
	"strings"
	"os"
)

const (
	VIRTIO_PORT_PATH = "/dev/virtio-ports/applianceVm.vport"
	BOOTSTRAP_INFO_CACHE = "/home/vyos/zvr/bootstrap-info.json"
)

type nic struct {
	mac string
	ip string
	name string
	netmask string
	isDefaultRoute bool
	gateway string
}

var bootstrapInfo map[string]interface{} = make(map[string]interface{})
var nics map[string]*nic = make(map[string]*nic)

func waitIptablesServiceOnline()  {
	bash := utils.Bash{
		Command: "/sbin/iptables-save",
	}

	utils.LoopRunUntilSuccessOrTimeout(func() bool {
		err := bash.Run()
		if err != nil {
			log.Debugf("iptables service seems not ready, %v", err)
		}
		return err == nil
	}, time.Duration(120)*time.Second, time.Duration(500)*time.Millisecond)
}

func waitVirtioPortOnline() {
	utils.LoopRunUntilSuccessOrTimeout(func() bool {
		ok, err := utils.PathExists(VIRTIO_PORT_PATH); utils.PanicOnError(err)
		if !ok {
			log.Debugf("%s doesn't not exist, wait it ...", VIRTIO_PORT_PATH)
		}
		return ok
	}, time.Duration(120)*time.Second, time.Duration(500)*time.Millisecond)
}

func parseBootInfo() {
	utils.LoopRunUntilSuccessOrTimeout(func() bool {
		content, err := ioutil.ReadFile(VIRTIO_PORT_PATH); utils.PanicOnError(err)
		if len(content) == 0 {
			log.Debugf("no content in %s, it may not be ready, wait it ...", VIRTIO_PORT_PATH)
			return false
		}

		if err := json.Unmarshal(content, &bootstrapInfo); err != nil {
			panic(errors.Wrap(err, fmt.Sprintf("unable to JSON parse:\n %s", string(content))))
		}

		err = utils.MkdirForFile(BOOTSTRAP_INFO_CACHE, 0666); utils.PanicOnError(err)
		err = ioutil.WriteFile(BOOTSTRAP_INFO_CACHE, content, 0666); utils.PanicOnError(err)
		err = os.Chmod(BOOTSTRAP_INFO_CACHE, 0777); utils.PanicOnError(err)
		log.Debugf("recieved bootstrap info:\n%s", string(content))
		return true
	}, time.Duration(300)*time.Second, time.Duration(1)*time.Second)
}

func configureVyos()  {
	mgmtNic := bootstrapInfo["managementNic"].(map[string]interface{})
	if mgmtNic == nil {
		panic(errors.New("no field 'managementNic' in bootstrap info"))
	}

	eth0 := &nic{ name: "eth0" }
	var ok bool
	eth0.mac, ok = mgmtNic["mac"].(string); utils.PanicIfError(ok, errors.New("cannot find 'mac' field for the management nic"))
	eth0.netmask, ok = mgmtNic["netmask"].(string); utils.PanicIfError(ok, errors.New("cannot find 'netmask' field for the management nic"))
	eth0.ip, ok = mgmtNic["ip"].(string); utils.PanicIfError(ok, errors.New("cannot find 'ip' field for the management nic"))
	eth0.isDefaultRoute = mgmtNic["isDefaultRoute"].(bool)
	eth0.gateway = mgmtNic["gateway"].(string)
	nics[eth0.name] = eth0

	otherNics := bootstrapInfo["additionalNics"].([]interface{})
	if otherNics != nil {
		for _, o := range otherNics {
			onic := o.(map[string]interface{})
			n := &nic{}
			n.name, ok = onic["deviceName"].(string); utils.PanicIfError(ok, fmt.Errorf("cannot find 'deviceName' field for the nic"))
			n.mac, ok = onic["mac"].(string); utils.PanicIfError(ok, errors.New("cannot find 'mac' field for the nic"))
			n.netmask, ok = onic["netmask"].(string); utils.PanicIfError(ok, fmt.Errorf("cannot find 'netmask' field for the nic[name:%s]", n.name))
			n.ip, ok = onic["ip"].(string); utils.PanicIfError(ok, fmt.Errorf("cannot find 'ip' field for the nic[name:%s]", n.name))
			n.gateway = onic["gateway"].(string)
			n.isDefaultRoute = onic["isDefaultRoute"].(bool)
			nics[n.name] = n
		}
	}

	type deviceName struct {
		expected string
		actual string
		swap string
	}

	devNames := make([]*deviceName, 0)

	// check integrity of nics
	for _, nic := range nics {
		utils.Assertf(nic.name != "", "name cannot be empty[mac:%s]", nic.mac)
		utils.Assertf(nic.ip != "", "ip cannot be empty[nicname: %s]", nic.name)
		utils.Assertf(nic.gateway != "", "gateway cannot be empty[nicname:%s]", nic.name)
		utils.Assertf(nic.netmask != "", "netmask cannot be empty[nicname:%s]", nic.name)
		utils.Assertf(nic.mac != "", "mac cannot be empty[nicname:%s]", nic.name)

		nicname, err := utils.GetNicNameByMac(nic.mac); utils.PanicOnError(err)
		if nicname != nic.name {
			devNames = append(devNames, &deviceName{
				expected: nic.name,
				actual: nicname,
			})
		}
	}

	if len(devNames) != 0 {
		// shutdown links and change to temporary names
		cmds := make([]string, 0)
		for i, devname := range devNames {
			devnum := 1000 + i

			devname.swap = fmt.Sprintf("eth%v", devnum)
			cmds = append(cmds, fmt.Sprintf("ip link set dev %v down", devname.actual))
			cmds = append(cmds, fmt.Sprintf("ip link set dev %v name %v", devname.actual, devname.swap))
		}

		b := utils.Bash{
			Command: strings.Join(cmds, "\n"),
		}

		b.Run()
		b.PanicIfError()

		// change temporary names to real names and bring up links
		cmds = make([]string, 0)
		for _, devname := range devNames {
			cmds = append(cmds, fmt.Sprintf("ip link set dev %v name %v", devname.swap, devname.expected))
			cmds = append(cmds, fmt.Sprintf("ip link set dev %v up", devname.expected))
		}

		b = utils.Bash{
			Command: strings.Join(cmds, "\n"),
		}

		b.Run()
		b.PanicIfError()
	}

	vyos := server.NewParserFromShowConfiguration()
	tree := vyos.Tree

	sshkey := bootstrapInfo["publicKey"].(string)
	utils.Assert(sshkey != "", "cannot find 'publicKey' in bootstrap info")
	sshkeyparts := strings.Split(sshkey, " ")
	sshtype := sshkeyparts[0]
	key := sshkeyparts[1]
	id := sshkeyparts[2]

	tree.Setf("system login user vyos authentication public-keys %s key %s", id, key)
	tree.Setf("system login user vyos authentication public-keys %s type %s", id, sshtype)

	setNic := func(nic *nic) {
		cidr, err := utils.NetmaskToCIDR(nic.netmask); utils.PanicOnError(err)
		//tree.Setf("interfaces ethernet %s hw-id %s", nic.name, nic.mac)
		tree.Setf("interfaces ethernet %s address %s", nic.name, fmt.Sprintf("%v/%v", nic.ip, cidr))
		tree.Setf("interfaces ethernet %s duplex auto", nic.name)
		tree.Setf("interfaces ethernet %s smp_affinity auto", nic.name)
		tree.Setf("interfaces ethernet %s speed auto", nic.name)
		if nic.isDefaultRoute {
			tree.Setf("system gateway-address %v", nic.gateway)
		}
	}

	sshport := bootstrapInfo["sshPort"].(float64)
	utils.Assert(sshport != 0, "sshport not found in bootstrap info")
	tree.Setf("service ssh port %v", int(sshport))
	tree.Setf("service ssh listen-address %v", eth0.ip)

	// configure firewall
	for _, nic := range nics {
		setNic(nic)

		tree.SetFirewallOnInterface(nic.name, "local",
			"action accept",
			"state established enable",
			"state related enable",
			fmt.Sprintf("destination address %v", nic.ip),
		)
		tree.SetFirewallOnInterface(nic.name, "local",
			"action accept",
			"protocol icmp",
			fmt.Sprintf("destination address %v", nic.ip),
		)

		if !nic.isDefaultRoute && nic.name != "eth0" {
			// for nic connecting to the guest network
			// make its FORWARDING chain open
			// here in = "-A FORWARDING -i nic.name -j ACCEPT"
			tree.SetFirewallOnInterface(nic.name, "in",
				"action accept",
				"state established enable",
				"state related enable",
				"state new enable",
			)
		} else {
			tree.SetFirewallOnInterface(nic.name, "in",
				"action accept",
				"state established enable",
				"state related enable",
			)
		}

		tree.SetFirewallOnInterface(nic.name, "in",
			"action accept",
			"protocol icmp",
		)

		// only allow ssh traffic on eth0, disable on others
		if nic.name == "eth0" {
			tree.SetFirewallOnInterface(nic.name, "local",
				fmt.Sprintf("destination port %v", int(sshport)),
				fmt.Sprintf("destination address %v", nic.ip),
				"protocol tcp",
				"action accept",
			)
		} else {
			tree.SetFirewallOnInterface(nic.name, "local",
				fmt.Sprintf("destination port %v", int(sshport)),
				fmt.Sprintf("destination address %v", nic.ip),
				"protocol tcp",
				"action reject",
			)
		}

		tree.SetFirewallDefaultAction(nic.name, "local", "reject")
		tree.SetFirewallDefaultAction(nic.name, "in", "reject")

		tree.AttachFirewallToInterface(nic.name, "local")
		tree.AttachFirewallToInterface(nic.name, "in")
	}

	tree.Set("system time-zone Asia/Shanghai")

	password := bootstrapInfo["vyosPassword"]; utils.Assert(password != "", "vyosPassword cannot be empty")
	tree.Setf("system login user vyos authentication plaintext-password %v", password)

	tree.Apply(true)

	arping := func(nicname, ip, gateway string) {
		b := utils.Bash{ Command: fmt.Sprintf("arping -A -U -c 1 -I %s -s %s %s", nicname, ip, gateway) }
		b.Run()
	}

	// arping to advocate our mac addresses
	arping("eth0", eth0.ip, eth0.gateway)
	for _, nic := range nics {
		arping(nic.name, nic.ip, nic.gateway)
	}
}

func startZvr()  {
	b := utils.Bash{
		Command: "/etc/init.d/zstack-virtualrouteragent restart",
	}
	b.Run()
	b.PanicIfError()
}

func main() {
	utils.InitLog("/home/vyos/zvr/zvrboot.log", false)
	waitIptablesServiceOnline()
	waitVirtioPortOnline()
	parseBootInfo()
	configureVyos()
	startZvr()
	log.Debugf("successfully configured the sysmtem and bootstrap the zstack virtual router agents")
}
<commit_msg>Add support to receive bootstrap info as ESX vm<commit_after>package main

import (
	"zvr/utils"
	"time"
	log "github.com/Sirupsen/logrus"
	"io/ioutil"
	"encoding/json"
	"github.com/pkg/errors"
	"fmt"
	"zvr/server"
	"strings"
	"os"
)

const (
	VIRTIO_PORT_PATH = "/dev/virtio-ports/applianceVm.vport"
	BOOTSTRAP_INFO_CACHE = "/home/vyos/zvr/bootstrap-info.json"
	TMP_LOCATION_FOR_ESX = "/tmp/bootstrap-info.json"
)

type nic struct {
	mac string
	ip string
	name string
	netmask string
	isDefaultRoute bool
	gateway string
}

var bootstrapInfo map[string]interface{} = make(map[string]interface{})
var nics map[string]*nic = make(map[string]*nic)

func waitIptablesServiceOnline()  {
	bash := utils.Bash{
		Command: "/sbin/iptables-save",
	}

	utils.LoopRunUntilSuccessOrTimeout(func() bool {
		err := bash.Run()
		if err != nil {
			log.Debugf("iptables service seems not ready, %v", err)
		}
		return err == nil
	}, time.Duration(120)*time.Second, time.Duration(500)*time.Millisecond)
}

func waitVirtioPortOnline() {
	utils.LoopRunUntilSuccessOrTimeout(func() bool {
		ok, err := utils.PathExists(VIRTIO_PORT_PATH); utils.PanicOnError(err)
		if !ok {
			log.Debugf("%s doesn't not exist, wait it ...", VIRTIO_PORT_PATH)
		}
		return ok
	}, time.Duration(120)*time.Second, time.Duration(500)*time.Millisecond)
}

func isOnVMwareHypervisor() bool {
	bash := utils.Bash{
		Command: "dmesg | grep -q 'Hypervisor.*VMware'",
	}

	if ret, _, _, err := bash.RunWithReturn(); ret == 0 && err == nil {
		return true
	}

	return false
}

func parseEsxBootInfo() {
	utils.LoopRunUntilSuccessOrTimeout(func() bool {
		if _, err := os.Stat(TMP_LOCATION_FOR_ESX); os.IsNotExist(err) {
			log.Debugf("bootstrap info not ready, waiting ...")
			return false
		}

		err := utils.MkdirForFile(BOOTSTRAP_INFO_CACHE, 0666); utils.PanicOnError(err)
		err = os.Rename(TMP_LOCATION_FOR_ESX, BOOTSTRAP_INFO_CACHE); utils.PanicOnError(err)
		err = os.Chmod(BOOTSTRAP_INFO_CACHE, 0777); utils.PanicOnError(err)
		content, _ := ioutil.ReadFile(BOOTSTRAP_INFO_CACHE)
		log.Debugf("recieved bootstrap info:\n%s", string(content))
		return true
	}, time.Duration(300)*time.Second, time.Duration(1)*time.Second)
}

func parseKvmBootInfo() {
	utils.LoopRunUntilSuccessOrTimeout(func() bool {
		content, err := ioutil.ReadFile(VIRTIO_PORT_PATH); utils.PanicOnError(err)
		if len(content) == 0 {
			log.Debugf("no content in %s, it may not be ready, wait it ...", VIRTIO_PORT_PATH)
			return false
		}

		if err := json.Unmarshal(content, &bootstrapInfo); err != nil {
			panic(errors.Wrap(err, fmt.Sprintf("unable to JSON parse:\n %s", string(content))))
		}

		err = utils.MkdirForFile(BOOTSTRAP_INFO_CACHE, 0666); utils.PanicOnError(err)
		err = ioutil.WriteFile(BOOTSTRAP_INFO_CACHE, content, 0666); utils.PanicOnError(err)
		err = os.Chmod(BOOTSTRAP_INFO_CACHE, 0777); utils.PanicOnError(err)
		log.Debugf("recieved bootstrap info:\n%s", string(content))
		return true
	}, time.Duration(300)*time.Second, time.Duration(1)*time.Second)
}

func configureVyos()  {
	mgmtNic := bootstrapInfo["managementNic"].(map[string]interface{})
	if mgmtNic == nil {
		panic(errors.New("no field 'managementNic' in bootstrap info"))
	}

	eth0 := &nic{ name: "eth0" }
	var ok bool
	eth0.mac, ok = mgmtNic["mac"].(string); utils.PanicIfError(ok, errors.New("cannot find 'mac' field for the management nic"))
	eth0.netmask, ok = mgmtNic["netmask"].(string); utils.PanicIfError(ok, errors.New("cannot find 'netmask' field for the management nic"))
	eth0.ip, ok = mgmtNic["ip"].(string); utils.PanicIfError(ok, errors.New("cannot find 'ip' field for the management nic"))
	eth0.isDefaultRoute = mgmtNic["isDefaultRoute"].(bool)
	eth0.gateway = mgmtNic["gateway"].(string)
	nics[eth0.name] = eth0

	otherNics := bootstrapInfo["additionalNics"].([]interface{})
	if otherNics != nil {
		for _, o := range otherNics {
			onic := o.(map[string]interface{})
			n := &nic{}
			n.name, ok = onic["deviceName"].(string); utils.PanicIfError(ok, fmt.Errorf("cannot find 'deviceName' field for the nic"))
			n.mac, ok = onic["mac"].(string); utils.PanicIfError(ok, errors.New("cannot find 'mac' field for the nic"))
			n.netmask, ok = onic["netmask"].(string); utils.PanicIfError(ok, fmt.Errorf("cannot find 'netmask' field for the nic[name:%s]", n.name))
			n.ip, ok = onic["ip"].(string); utils.PanicIfError(ok, fmt.Errorf("cannot find 'ip' field for the nic[name:%s]", n.name))
			n.gateway = onic["gateway"].(string)
			n.isDefaultRoute = onic["isDefaultRoute"].(bool)
			nics[n.name] = n
		}
	}

	type deviceName struct {
		expected string
		actual string
		swap string
	}

	devNames := make([]*deviceName, 0)

	// check integrity of nics
	for _, nic := range nics {
		utils.Assertf(nic.name != "", "name cannot be empty[mac:%s]", nic.mac)
		utils.Assertf(nic.ip != "", "ip cannot be empty[nicname: %s]", nic.name)
		utils.Assertf(nic.gateway != "", "gateway cannot be empty[nicname:%s]", nic.name)
		utils.Assertf(nic.netmask != "", "netmask cannot be empty[nicname:%s]", nic.name)
		utils.Assertf(nic.mac != "", "mac cannot be empty[nicname:%s]", nic.name)

		nicname, err := utils.GetNicNameByMac(nic.mac); utils.PanicOnError(err)
		if nicname != nic.name {
			devNames = append(devNames, &deviceName{
				expected: nic.name,
				actual: nicname,
			})
		}
	}

	if len(devNames) != 0 {
		// shutdown links and change to temporary names
		cmds := make([]string, 0)
		for i, devname := range devNames {
			devnum := 1000 + i

			devname.swap = fmt.Sprintf("eth%v", devnum)
			cmds = append(cmds, fmt.Sprintf("ip link set dev %v down", devname.actual))
			cmds = append(cmds, fmt.Sprintf("ip link set dev %v name %v", devname.actual, devname.swap))
		}

		b := utils.Bash{
			Command: strings.Join(cmds, "\n"),
		}

		b.Run()
		b.PanicIfError()

		// change temporary names to real names and bring up links
		cmds = make([]string, 0)
		for _, devname := range devNames {
			cmds = append(cmds, fmt.Sprintf("ip link set dev %v name %v", devname.swap, devname.expected))
			cmds = append(cmds, fmt.Sprintf("ip link set dev %v up", devname.expected))
		}

		b = utils.Bash{
			Command: strings.Join(cmds, "\n"),
		}

		b.Run()
		b.PanicIfError()
	}

	vyos := server.NewParserFromShowConfiguration()
	tree := vyos.Tree

	sshkey := bootstrapInfo["publicKey"].(string)
	utils.Assert(sshkey != "", "cannot find 'publicKey' in bootstrap info")
	sshkeyparts := strings.Split(sshkey, " ")
	sshtype := sshkeyparts[0]
	key := sshkeyparts[1]
	id := sshkeyparts[2]

	tree.Setf("system login user vyos authentication public-keys %s key %s", id, key)
	tree.Setf("system login user vyos authentication public-keys %s type %s", id, sshtype)

	setNic := func(nic *nic) {
		cidr, err := utils.NetmaskToCIDR(nic.netmask); utils.PanicOnError(err)
		//tree.Setf("interfaces ethernet %s hw-id %s", nic.name, nic.mac)
		tree.Setf("interfaces ethernet %s address %s", nic.name, fmt.Sprintf("%v/%v", nic.ip, cidr))
		tree.Setf("interfaces ethernet %s duplex auto", nic.name)
		tree.Setf("interfaces ethernet %s smp_affinity auto", nic.name)
		tree.Setf("interfaces ethernet %s speed auto", nic.name)
		if nic.isDefaultRoute {
			tree.Setf("system gateway-address %v", nic.gateway)
		}
	}

	sshport := bootstrapInfo["sshPort"].(float64)
	utils.Assert(sshport != 0, "sshport not found in bootstrap info")
	tree.Setf("service ssh port %v", int(sshport))
	tree.Setf("service ssh listen-address %v", eth0.ip)

	// configure firewall
	for _, nic := range nics {
		setNic(nic)

		tree.SetFirewallOnInterface(nic.name, "local",
			"action accept",
			"state established enable",
			"state related enable",
			fmt.Sprintf("destination address %v", nic.ip),
		)
		tree.SetFirewallOnInterface(nic.name, "local",
			"action accept",
			"protocol icmp",
			fmt.Sprintf("destination address %v", nic.ip),
		)

		if !nic.isDefaultRoute && nic.name != "eth0" {
			// for nic connecting to the guest network
			// make its FORWARDING chain open
			// here in = "-A FORWARDING -i nic.name -j ACCEPT"
			tree.SetFirewallOnInterface(nic.name, "in",
				"action accept",
				"state established enable",
				"state related enable",
				"state new enable",
			)
		} else {
			tree.SetFirewallOnInterface(nic.name, "in",
				"action accept",
				"state established enable",
				"state related enable",
			)
		}

		tree.SetFirewallOnInterface(nic.name, "in",
			"action accept",
			"protocol icmp",
		)

		// only allow ssh traffic on eth0, disable on others
		if nic.name == "eth0" {
			tree.SetFirewallOnInterface(nic.name, "local",
				fmt.Sprintf("destination port %v", int(sshport)),
				fmt.Sprintf("destination address %v", nic.ip),
				"protocol tcp",
				"action accept",
			)
		} else {
			tree.SetFirewallOnInterface(nic.name, "local",
				fmt.Sprintf("destination port %v", int(sshport)),
				fmt.Sprintf("destination address %v", nic.ip),
				"protocol tcp",
				"action reject",
			)
		}

		tree.SetFirewallDefaultAction(nic.name, "local", "reject")
		tree.SetFirewallDefaultAction(nic.name, "in", "reject")

		tree.AttachFirewallToInterface(nic.name, "local")
		tree.AttachFirewallToInterface(nic.name, "in")
	}

	tree.Set("system time-zone Asia/Shanghai")

	password := bootstrapInfo["vyosPassword"]; utils.Assert(password != "", "vyosPassword cannot be empty")
	tree.Setf("system login user vyos authentication plaintext-password %v", password)

	tree.Apply(true)

	arping := func(nicname, ip, gateway string) {
		b := utils.Bash{ Command: fmt.Sprintf("arping -A -U -c 1 -I %s -s %s %s", nicname, ip, gateway) }
		b.Run()
	}

	// arping to advocate our mac addresses
	arping("eth0", eth0.ip, eth0.gateway)
	for _, nic := range nics {
		arping(nic.name, nic.ip, nic.gateway)
	}
}

func startZvr()  {
	b := utils.Bash{
		Command: "/etc/init.d/zstack-virtualrouteragent restart",
	}
	b.Run()
	b.PanicIfError()
}

func main() {
	utils.InitLog("/home/vyos/zvr/zvrboot.log", false)
	waitIptablesServiceOnline()
	waitVirtioPortOnline()
	if isOnVMwareHypervisor() {
		parseEsxBootInfo()
	} else {
		parseKvmBootInfo()
	}
	configureVyos()
	startZvr()
	log.Debugf("successfully configured the sysmtem and bootstrap the zstack virtual router agents")
}
<|endoftext|>


<commit_before>package mesh

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"strings"
	"time"

	"github.com/emersion/go-upnp-igd"
)

var (
	ErrAllAddressesFailed = errors.New("all addresses failed to dial")
)

type Net struct {
	registry Registry
	accepted chan net.Conn
	reusable map[string]*reusableConn
	handlers map[string]Handler
	close    chan bool
}

func New(registry Registry) *Net {
	return &Net{
		registry: registry,
		close:    make(chan bool),
		accepted: make(chan net.Conn),
		reusable: map[string]*reusableConn{},
		handlers: map[string]Handler{
			"id":    &ID{},
			"yamux": &Yamux{},
		},
	}
}

func (n *Net) RegisterHandler(protocol string, handler Handler) error {
	n.handlers[protocol] = handler
	return nil
}

func (n *Net) Dial(ctx context.Context, peerID string, commands ...string) (net.Conn, error) {
	peerInfo, err := n.registry.GetPeerInfo(peerID)
	if err != nil {
		return nil, err
	}
	var conn net.Conn
	if reusableConn, ok := n.reusable[peerID]; ok {
		newConn, err := reusableConn.NewConn()
		if err != nil {
			// TODO remove reusable conn
		} else {
			fmt.Println("dial reusing conn")
			conn = newConn
		}
	}
	if conn == nil {
		for _, addr := range peerInfo.Addresses {
			addr = strings.Replace(addr, "tcp:", "", 1)
			fmt.Println("dial dialing new conn to", addr)
			dialer := net.Dialer{Timeout: time.Second * 5}
			newConn, err := dialer.DialContext(ctx, "tcp", addr)
			if err != nil {
				// TODO blacklist address for a bit
				// TODO hold error maybe?
				// return nil, err
				continue
			}
			localAddress := peerAddress{
				network: "tcp",
				peerID:  n.registry.GetLocalPeerInfo().ID,
			}
			remoteAddress := peerAddress{
				network: "tcp",
				peerID:  peerID,
			}
			conn = NewAddressableConn(newConn, localAddress, remoteAddress)
			break
		}
	}
	if conn == nil {
		return nil, ErrAllAddressesFailed
	}

	finalConn, err := n.Select(conn, commands...)
	if err != nil {
		if err := conn.Close(); err != nil {
			fmt.Println("could not close connection after failure to select")
		}
		return nil, err
	}

	return finalConn, nil
}

func (n *Net) Select(conn net.Conn, commands ...string) (net.Conn, error) {
	for _, command := range commands {
		if err := WriteToken(conn, []byte(command)); err != nil {
			return nil, err
		}
		token, err := ReadToken(conn)
		if err != nil {
			return nil, err
		}
		fmt.Printf("Dialer got token %s for command %s\n", string(token), command)
		if string(token) != "ok" {
			return nil, errors.New("unexpected token response")
		}
		handler, ok := n.handlers[command]
		if !ok {
			return nil, errors.New("no such handler")
		}
		newConn, err := handler.Initiate(conn)
		if err != nil {
			return nil, err
		}
		// TODO maybe use a switch
		if reusableConn, ok := newConn.(*reusableConn); ok {
			go reusableConn.Accepted(n.accepted)
			// TODO lock
			fmt.Println("client storing reusable")
			n.reusable[conn.RemoteAddr().String()] = reusableConn
		}
		conn = newConn
	}
	return conn, nil
}

// TODO do we need to return a listener?
func (n *Net) Listen(addr string) (Listener, string, error) {
	tcpListener, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, "", err
	}
	port := tcpListener.Addr().(*net.TCPAddr).Port
	addresses := GetAddresses(tcpListener)

	devices := make(chan igd.Device)
	go func() {
		for device := range devices {
			externalAddress, err := device.GetExternalIPAddress()
			if err != nil {
				fmt.Println("could not get external ip")
				continue
			}
			desc := "nimona"
			ttl := time.Hour * 24 * 365
			if _, err := device.AddPortMapping(igd.TCP, port, port, desc, ttl); err != nil {
				fmt.Println("could not add port mapping", err)
			} else {
				addresses = append(addresses, fmt.Sprintf("tcp:%s:%d", externalAddress.String(), port))
			}
		}
	}()

	if err := igd.Discover(devices, 5*time.Second); err != nil {
		log.Println("could not discover devices")
	}

	n.registry.PutLocalPeerInfo(&PeerInfo{
		ID:        n.registry.GetLocalPeerInfo().ID,
		Addresses: addresses,
	})

	go func() {
		for {
			conn := <-n.accepted
			go n.HandleSelection(conn)
		}
	}()

	closed := false

	go func() {
		closed = true
		<-n.close
		fmt.Println("Closing")
		tcpListener.Close()
	}()

	go func() {
		for {
			conn, err := tcpListener.Accept()
			if err != nil {
				if closed {
					return
				}
				fmt.Println("Error accepting: ", err.Error())
				// TODO check conn is still alive and return
				return
			}
			localAddress := peerAddress{
				network: "tcp",
				peerID:  n.registry.GetLocalPeerInfo().ID,
			}
			remoteAddress := peerAddress{
				network: "tcp",
				peerID:  conn.RemoteAddr().String(),
			}
			conn = NewAddressableConn(conn, localAddress, remoteAddress)
			n.accepted <- conn
		}
	}()

	return nil, tcpListener.Addr().String(), nil
}

func (n *Net) Close() error {
	n.close <- true
	return nil
}

func (n *Net) HandleSelection(conn net.Conn) (net.Conn, error) {
	for {
		token, err := ReadToken(conn)
		if err != nil {
			return nil, err
		}
		fmt.Println("selection handler got token", string(token))
		handler, ok := n.handlers[string(token)]
		if !ok {
			if err := WriteToken(conn, []byte("error")); err != nil {
				return nil, err
			}
			return nil, errors.New("no such handler")
		}
		if err := WriteToken(conn, []byte("ok")); err != nil {
			return nil, err
		}
		newConn, err := handler.Handle(conn)
		if err != nil {
			return nil, err
		}
		// TODO maybe use a switch
		if reusableConn, ok := newConn.(*reusableConn); ok {
			go reusableConn.Accepted(n.accepted)
			// TODO lock
			fmt.Println("server storing reusable")
			n.reusable[conn.RemoteAddr().String()] = reusableConn
		}
		// TODO missing peerID to reuse conn
		// TODO maybe use a switch
		// if reusableConn, ok := newConn.(ReusableConn); ok {
		// 	// TODO lock
		// 	n.reusable[peerID] = reusableConn
		// }
		conn = newConn
	}
}
<commit_msg>Force id and yamux wrappers<commit_after>package mesh

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"strings"
	"time"

	"github.com/emersion/go-upnp-igd"
)

var (
	ErrAllAddressesFailed = errors.New("all addresses failed to dial")
)

type Net struct {
	registry Registry
	accepted chan net.Conn
	reusable map[string]*reusableConn
	handlers map[string]Handler
	close    chan bool
}

func New(registry Registry) *Net {
	return &Net{
		registry: registry,
		close:    make(chan bool),
		accepted: make(chan net.Conn),
		reusable: map[string]*reusableConn{},
		handlers: map[string]Handler{
			"id":    &ID{},
			"yamux": &Yamux{},
		},
	}
}

func (n *Net) RegisterHandler(protocol string, handler Handler) error {
	n.handlers[protocol] = handler
	return nil
}

func (n *Net) Dial(ctx context.Context, peerID string, commands ...string) (net.Conn, error) {
	peerInfo, err := n.registry.GetPeerInfo(peerID)
	if err != nil {
		return nil, err
	}
	var conn net.Conn
	if reusableConn, ok := n.reusable[peerID]; ok {
		newConn, err := reusableConn.NewConn()
		if err != nil {
			// TODO remove reusable conn
		} else {
			fmt.Println("dial reusing conn")
			conn = newConn
		}
	}
	if conn == nil {
		for _, addr := range peerInfo.Addresses {
			addr = strings.Replace(addr, "tcp:", "", 1)
			fmt.Println("dial dialing new conn to", addr)
			dialer := net.Dialer{Timeout: time.Second * 5}
			newConn, err := dialer.DialContext(ctx, "tcp", addr)
			if err != nil {
				// TODO blacklist address for a bit
				// TODO hold error maybe?
				// return nil, err
				continue
			}
			localAddress := peerAddress{
				network: "tcp",
				peerID:  n.registry.GetLocalPeerInfo().ID,
			}
			remoteAddress := peerAddress{
				network: "tcp",
				peerID:  peerID,
			}
			conn = NewAddressableConn(newConn, localAddress, remoteAddress)
			break
		}
	}
	if conn == nil {
		return nil, ErrAllAddressesFailed
	}

	commands = append([]string{"id", "yamux"}, commands...)
	finalConn, err := n.Select(conn, commands...)
	if err != nil {
		if err := conn.Close(); err != nil {
			fmt.Println("could not close connection after failure to select")
		}
		return nil, err
	}

	return finalConn, nil
}

func (n *Net) Select(conn net.Conn, commands ...string) (net.Conn, error) {
	for _, command := range commands {
		if err := WriteToken(conn, []byte(command)); err != nil {
			return nil, err
		}
		token, err := ReadToken(conn)
		if err != nil {
			return nil, err
		}
		fmt.Printf("Dialer got token %s for command %s\n", string(token), command)
		if string(token) != "ok" {
			return nil, errors.New("unexpected token response")
		}
		handler, ok := n.handlers[command]
		if !ok {
			return nil, errors.New("no such handler")
		}
		newConn, err := handler.Initiate(conn)
		if err != nil {
			return nil, err
		}
		// TODO maybe use a switch
		if reusableConn, ok := newConn.(*reusableConn); ok {
			go reusableConn.Accepted(n.accepted)
			// TODO lock
			fmt.Println("client storing reusable")
			n.reusable[conn.RemoteAddr().String()] = reusableConn
		}
		conn = newConn
	}
	return conn, nil
}

// TODO do we need to return a listener?
func (n *Net) Listen(addr string) (Listener, string, error) {
	tcpListener, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, "", err
	}
	port := tcpListener.Addr().(*net.TCPAddr).Port
	addresses := GetAddresses(tcpListener)

	devices := make(chan igd.Device)
	go func() {
		for device := range devices {
			externalAddress, err := device.GetExternalIPAddress()
			if err != nil {
				fmt.Println("could not get external ip")
				continue
			}
			desc := "nimona"
			ttl := time.Hour * 24 * 365
			if _, err := device.AddPortMapping(igd.TCP, port, port, desc, ttl); err != nil {
				fmt.Println("could not add port mapping", err)
			} else {
				addresses = append(addresses, fmt.Sprintf("tcp:%s:%d", externalAddress.String(), port))
			}
		}
	}()

	if err := igd.Discover(devices, 5*time.Second); err != nil {
		log.Println("could not discover devices")
	}

	n.registry.PutLocalPeerInfo(&PeerInfo{
		ID:        n.registry.GetLocalPeerInfo().ID,
		Addresses: addresses,
	})

	go func() {
		for {
			conn := <-n.accepted
			go n.HandleSelection(conn)
		}
	}()

	closed := false

	go func() {
		closed = true
		<-n.close
		fmt.Println("Closing")
		tcpListener.Close()
	}()

	go func() {
		for {
			conn, err := tcpListener.Accept()
			if err != nil {
				if closed {
					return
				}
				fmt.Println("Error accepting: ", err.Error())
				// TODO check conn is still alive and return
				return
			}
			localAddress := peerAddress{
				network: "tcp",
				peerID:  n.registry.GetLocalPeerInfo().ID,
			}
			remoteAddress := peerAddress{
				network: "tcp",
				peerID:  conn.RemoteAddr().String(),
			}
			conn = NewAddressableConn(conn, localAddress, remoteAddress)
			n.accepted <- conn
		}
	}()

	return nil, tcpListener.Addr().String(), nil
}

func (n *Net) Close() error {
	n.close <- true
	return nil
}

func (n *Net) HandleSelection(conn net.Conn) (net.Conn, error) {
	for {
		token, err := ReadToken(conn)
		if err != nil {
			return nil, err
		}
		fmt.Println("selection handler got token", string(token))
		handler, ok := n.handlers[string(token)]
		if !ok {
			if err := WriteToken(conn, []byte("error")); err != nil {
				return nil, err
			}
			return nil, errors.New("no such handler")
		}
		if err := WriteToken(conn, []byte("ok")); err != nil {
			return nil, err
		}
		newConn, err := handler.Handle(conn)
		if err != nil {
			return nil, err
		}
		// TODO maybe use a switch
		if reusableConn, ok := newConn.(*reusableConn); ok {
			go reusableConn.Accepted(n.accepted)
			// TODO lock
			fmt.Println("server storing reusable")
			n.reusable[conn.RemoteAddr().String()] = reusableConn
		}
		conn = newConn
	}
}
<|endoftext|>