<commit_before>package gce

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/emccode/rexray/core"
	"github.com/emccode/rexray/core/config"
	"github.com/emccode/rexray/core/errors"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/compute/v1"
	"io/ioutil"
	"regexp"
	"strconv"
	"strings"
)

const providerName = "gce"

// The GCE storage driver.
type driver struct {
	client  *compute.Service
	r       *core.RexRay
	zone    string
	project string
}

func ef() errors.Fields {
	return errors.Fields{
		"provider": providerName,
	}
}

func eff(fields errors.Fields) map[string]interface{} {
	errFields := map[string]interface{}{
		"provider": providerName,
	}
	if fields != nil {
		for k, v := range fields {
			errFields[k] = v
		}
	}
	return errFields
}

func init() {
	core.RegisterDriver(providerName, newDriver)
	config.Register(configRegistration())
}

func newDriver() core.Driver {
	return &driver{}
}

func (d *driver) Init(r *core.RexRay) error {
	d.r = r

	var err error

	d.zone = d.r.Config.GetString("gce.zone")
	d.project = d.r.Config.GetString("gce.project")
	serviceAccountJSON, err := ioutil.ReadFile(d.r.Config.GetString("gce.keyfile"))
	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not read service account credentials file, %s => {%s}", d.r.Config.GetString("gce.keyfile"), err)
	}

	config, err := google.JWTConfigFromJSON(serviceAccountJSON,
		compute.ComputeScope,
	)
	client, err := compute.New(config.Client(context.Background()))

	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not create compute client => {%s}", err)
	}
	d.client = client
	log.WithField("provider", providerName).Info("storage driver initialized")

	return nil
}

func (d *driver) Name() string {
	return providerName
}

func (d *driver) GetVolumeMapping() ([]*core.BlockDevice, error) {
	log.WithField("provider", providerName).Debug("GetVolumeMapping")
	return nil, nil
}

func (d *driver) GetInstance() (*core.Instance, error) {
	log.WithField("provider", providerName).Debug("GetInstance")
	return nil, nil
}

func (d *driver) CreateSnapshot(
	runAsync bool,
	snapshotName, volumeID, description string) ([]*core.Snapshot, error) {

	log.WithField("provider", providerName).Debug("CreateSnapshot")
	return nil, nil

}

func (d *driver) GetSnapshot(
	volumeID, snapshotID, snapshotName string) ([]*core.Snapshot, error) {

	log.WithField("provider", providerName).Debug("GetSnapshot")
	return nil, nil
}

func (d *driver) RemoveSnapshot(snapshotID string) error {
	log.WithField("provider", providerName).Debug("RemoveSnapshot")
	return nil
}

func (d *driver) GetDeviceNextAvailable() (string, error) {
	letters := []string{
		"a", "b", "c", "d", "e", "f", "g", "h",
		"i", "j", "k", "l", "m", "n", "o", "p"}

	blockDeviceNames := make(map[string]bool)

	blockDeviceMapping, err := d.GetVolumeMapping()
	if err != nil {
		return "", err
	}

	for _, blockDevice := range blockDeviceMapping {
		re, _ := regexp.Compile(`^/dev/xvd([a-z])`)
		res := re.FindStringSubmatch(blockDevice.DeviceName)
		if len(res) > 0 {
			blockDeviceNames[res[1]] = true
		}
	}

	localDevices, err := getLocalDevices()
	if err != nil {
		return "", err
	}

	for _, localDevice := range localDevices {
		re, _ := regexp.Compile(`^xvd([a-z])`)
		res := re.FindStringSubmatch(localDevice)
		if len(res) > 0 {
			blockDeviceNames[res[1]] = true
		}
	}

	for _, letter := range letters {
		if !blockDeviceNames[letter] {
			nextDeviceName := "/dev/xvd" + letter
			log.Println("Got next device name: " + nextDeviceName)
			return nextDeviceName, nil
		}
	}
	return "", errors.New("No available device")
}

func getLocalDevices() (deviceNames []string, err error) {
	file := "/proc/partitions"
	contentBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return []string{}, err
	}

	content := string(contentBytes)

	lines := strings.Split(content, "\n")
	for _, line := range lines[2:] {
		fields := strings.Fields(line)
		if len(fields) == 4 {
			deviceNames = append(deviceNames, fields[3])
		}
	}

	return deviceNames, nil
}

func (d *driver) CreateVolume(
	runAsync bool, volumeName, volumeID, snapshotID, volumeType string,
	IOPS, size int64, availabilityZone string) (*core.Volume, error) {
	log.WithField("provider", providerName).Debug("CreateVolume")
	return nil, nil

}

func (d *driver) createVolumeCreateSnapshot(
	volumeID string, snapshotID string) (string, error) {
	log.WithField("provider", providerName).Debug("CreateVolumeCreateSnapshot")
	return "", nil

}

func (d *driver) GetVolume(
	volumeID, volumeName string) ([]*core.Volume, error) {
	log.WithField("provider", providerName).Debugf("GetVolume :%s %s", volumeID, volumeName)

	query := d.client.Disks.List(d.project, d.zone)
	if volumeID != "" {
		query.Filter(fmt.Sprintf("id eq %s", volumeID))
	}
	if volumeName != "" {
		query.Filter(fmt.Sprintf("name eq %s", volumeName))
	}
	var attachments []*core.VolumeAttachment
	instances, err := d.client.Instances.List(d.project, d.zone).Do()
	for _, instance := range instances.Items {
		for _, disk := range instance.Disks {
			attachment := &core.VolumeAttachment{
				InstanceID: strconv.FormatUint(instance.Id, 10),
				DeviceName: disk.DeviceName,
				Status:     disk.Mode,
				VolumeID:   disk.Source,
			}
			attachments = append(attachments, attachment)

		}
	}

	disks, err := query.Do()
	if err != nil {
		return []*core.Volume{}, err
	}
	var volumesSD []*core.Volume
	for _, disk := range disks.Items {
		var diskAttachments []*core.VolumeAttachment
		for _, attachment := range attachments {
			if attachment.VolumeID == disk.SelfLink {
				diskAttachments = append(diskAttachments, &core.VolumeAttachment{
					InstanceID: attachment.InstanceID,
					DeviceName: attachment.DeviceName,
					Status:     attachment.Status,
					VolumeID:   strconv.FormatUint(disk.Id, 10),
				})
			}
		}
		volumeSD := &core.Volume{
			Name:             disk.Name,
			VolumeID:         strconv.FormatUint(disk.Id, 10),
			AvailabilityZone: disk.Zone,
			Status:           disk.Status,
			VolumeType:       disk.Kind,
			NetworkName:      disk.SelfLink,
			IOPS:             0,
			Size:             strconv.FormatInt(disk.SizeGb, 10),
			Attachments:      diskAttachments,
		}
		volumesSD = append(volumesSD, volumeSD)

	}
	return volumesSD, nil
}

func (d *driver) GetVolumeAttach(
	volumeID, instanceID string) ([]*core.VolumeAttachment, error) {
	log.WithField("provider", providerName).Debug("GetVolumeAttach")
	return nil, nil
}

func (d *driver) waitSnapshotComplete(snapshotID string) error {
	return nil
}

func (d *driver) waitVolumeComplete(volumeID string) error {
	return nil
}

func (d *driver) waitVolumeAttach(volumeID, instanceID string) error {
	return nil
}

func (d *driver) waitVolumeDetach(volumeID string) error {
	return nil
}

func (d *driver) RemoveVolume(volumeID string) error {
	return nil
}

func (d *driver) AttachVolume(
	runAsync bool,
	volumeID, instanceID string) ([]*core.VolumeAttachment, error) {
	log.WithField("provider", providerName).Debug("AttachVolume")
	return nil, nil

}

func (d *driver) DetachVolume(
	runAsync bool,
	volumeID, blank string) error {
	log.WithField("provider", providerName).Debug("DetachVolume")
	return nil
}

func (d *driver) CopySnapshot(runAsync bool,
	volumeID, snapshotID, snapshotName, destinationSnapshotName,
	destinationRegion string) (*core.Snapshot, error) {
	log.WithField("provider", providerName).Debug("CopySnapshot")
	return nil, nil
}

func configRegistration() *config.Registration {
	r := config.NewRegistration("Google GCE")
	r.Key(config.String, "", "", "", "gce.zone")
	r.Key(config.String, "", "", "", "gce.project")
	r.Key(config.String, "", "", "", "gce.keyfile")
	return r
}
<commit_msg>Implement map method<commit_after>package gce

import (
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/emccode/rexray/core"
	"github.com/emccode/rexray/core/config"
	"github.com/emccode/rexray/core/errors"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/compute/v1"
	"io/ioutil"
	"regexp"
	"strconv"
	"strings"
)

const providerName = "gce"

// The GCE storage driver.
type driver struct {
	client  *compute.Service
	r       *core.RexRay
	zone    string
	project string
}

func ef() errors.Fields {
	return errors.Fields{
		"provider": providerName,
	}
}

func eff(fields errors.Fields) map[string]interface{} {
	errFields := map[string]interface{}{
		"provider": providerName,
	}
	if fields != nil {
		for k, v := range fields {
			errFields[k] = v
		}
	}
	return errFields
}

func init() {
	core.RegisterDriver(providerName, newDriver)
	config.Register(configRegistration())
}

func newDriver() core.Driver {
	return &driver{}
}

func (d *driver) Init(r *core.RexRay) error {
	d.r = r

	var err error

	d.zone = d.r.Config.GetString("gce.zone")
	d.project = d.r.Config.GetString("gce.project")
	serviceAccountJSON, err := ioutil.ReadFile(d.r.Config.GetString("gce.keyfile"))
	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not read service account credentials file, %s => {%s}", d.r.Config.GetString("gce.keyfile"), err)
	}

	config, err := google.JWTConfigFromJSON(serviceAccountJSON,
		compute.ComputeScope,
	)
	client, err := compute.New(config.Client(context.Background()))

	if err != nil {
		log.WithField("provider", providerName).Fatalf("Could not create compute client => {%s}", err)
	}
	d.client = client
	log.WithField("provider", providerName).Info("storage driver initialized")

	return nil
}

func (d *driver) Name() string {
	return providerName
}

func (d *driver) GetVolumeMapping() ([]*core.BlockDevice, error) {
	log.WithField("provider", providerName).Debug("GetVolumeMapping")
	return nil, nil
}

func (d *driver) GetInstance() (*core.Instance, error) {
	log.WithField("provider", providerName).Debug("GetInstance")
	return nil, nil
}

func (d *driver) CreateSnapshot(
	runAsync bool,
	snapshotName, volumeID, description string) ([]*core.Snapshot, error) {

	log.WithField("provider", providerName).Debug("CreateSnapshot")
	return nil, nil

}

func (d *driver) GetSnapshot(
	volumeID, snapshotID, snapshotName string) ([]*core.Snapshot, error) {

	log.WithField("provider", providerName).Debug("GetSnapshot")
	return nil, nil
}

func (d *driver) RemoveSnapshot(snapshotID string) error {
	log.WithField("provider", providerName).Debug("RemoveSnapshot")
	return nil
}

func (d *driver) GetDeviceNextAvailable() (string, error) {
	letters := []string{
		"a", "b", "c", "d", "e", "f", "g", "h",
		"i", "j", "k", "l", "m", "n", "o", "p"}

	blockDeviceNames := make(map[string]bool)

	blockDeviceMapping, err := d.GetVolumeMapping()
	if err != nil {
		return "", err
	}

	for _, blockDevice := range blockDeviceMapping {
		re, _ := regexp.Compile(`^/dev/xvd([a-z])`)
		res := re.FindStringSubmatch(blockDevice.DeviceName)
		if len(res) > 0 {
			blockDeviceNames[res[1]] = true
		}
	}

	localDevices, err := getLocalDevices()
	if err != nil {
		return "", err
	}

	for _, localDevice := range localDevices {
		re, _ := regexp.Compile(`^xvd([a-z])`)
		res := re.FindStringSubmatch(localDevice)
		if len(res) > 0 {
			blockDeviceNames[res[1]] = true
		}
	}

	for _, letter := range letters {
		if !blockDeviceNames[letter] {
			nextDeviceName := "/dev/xvd" + letter
			log.Println("Got next device name: " + nextDeviceName)
			return nextDeviceName, nil
		}
	}
	return "", errors.New("No available device")
}

func getLocalDevices() (deviceNames []string, err error) {
	file := "/proc/partitions"
	contentBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return []string{}, err
	}

	content := string(contentBytes)

	lines := strings.Split(content, "\n")
	for _, line := range lines[2:] {
		fields := strings.Fields(line)
		if len(fields) == 4 {
			deviceNames = append(deviceNames, fields[3])
		}
	}

	return deviceNames, nil
}

func (d *driver) CreateVolume(
	runAsync bool, volumeName, volumeID, snapshotID, volumeType string,
	IOPS, size int64, availabilityZone string) (*core.Volume, error) {
	log.WithField("provider", providerName).Debug("CreateVolume")
	return nil, nil

}

func (d *driver) createVolumeCreateSnapshot(
	volumeID string, snapshotID string) (string, error) {
	log.WithField("provider", providerName).Debug("CreateVolumeCreateSnapshot")
	return "", nil

}

func (d *driver) GetVolume(
	volumeID, volumeName string) ([]*core.Volume, error) {
	log.WithField("provider", providerName).Debugf("GetVolume :%s %s", volumeID, volumeName)

	query := d.client.Disks.List(d.project, d.zone)
	if volumeID != "" {
		query.Filter(fmt.Sprintf("id eq %s", volumeID))
	}
	if volumeName != "" {
		query.Filter(fmt.Sprintf("name eq %s", volumeName))
	}
	var attachments []*core.VolumeAttachment
	instances, err := d.client.Instances.List(d.project, d.zone).Do()
	if err != nil {
		return []*core.Volume{}, err
	}
	for _, instance := range instances.Items {
		for _, disk := range instance.Disks {
			attachment := &core.VolumeAttachment{
				InstanceID: strconv.FormatUint(instance.Id, 10),
				DeviceName: disk.DeviceName,
				Status:     disk.Mode,
				VolumeID:   disk.Source,
			}
			attachments = append(attachments, attachment)

		}
	}

	disks, err := query.Do()
	if err != nil {
		return []*core.Volume{}, err
	}
	var volumesSD []*core.Volume
	for _, disk := range disks.Items {
		var diskAttachments []*core.VolumeAttachment
		for _, attachment := range attachments {
			if attachment.VolumeID == disk.SelfLink {
				diskAttachments = append(diskAttachments, &core.VolumeAttachment{
					InstanceID: attachment.InstanceID,
					DeviceName: attachment.DeviceName,
					Status:     attachment.Status,
					VolumeID:   strconv.FormatUint(disk.Id, 10),
				})
			}
		}
		volumeSD := &core.Volume{
			Name:             disk.Name,
			VolumeID:         strconv.FormatUint(disk.Id, 10),
			AvailabilityZone: disk.Zone,
			Status:           disk.Status,
			VolumeType:       disk.Kind,
			NetworkName:      disk.SelfLink,
			IOPS:             0,
			Size:             strconv.FormatInt(disk.SizeGb, 10),
			Attachments:      diskAttachments,
		}
		volumesSD = append(volumesSD, volumeSD)

	}
	return volumesSD, nil
}

func (d *driver) GetVolumeAttach(
	volumeID, instanceID string) ([]*core.VolumeAttachment, error) {
	log.WithField("provider", providerName).Debug("GetVolumeAttach")
	var attachments []*core.VolumeAttachment
	query := d.client.Instances.List(d.project, d.zone)
	if instanceID != "" {
		query.Filter(fmt.Sprintf("id eq %s", instanceID))
	}
	instances, err := query.Do()
	if err != nil {
		return []*core.Volume{}, err
	}
	for _, instance := range instances.Items {
		for _, disk := range instance.Disks {
			attachment := &core.VolumeAttachment{
				InstanceID: strconv.FormatUint(instance.Id, 10),
				DeviceName: disk.DeviceName,
				Status:     disk.Mode,
				VolumeID:   disk.Source,
			}
			attachments = append(attachments, attachment)

		}
	}

}

func (d *driver) waitSnapshotComplete(snapshotID string) error {
	return nil
}

func (d *driver) waitVolumeComplete(volumeID string) error {
	return nil
}

func (d *driver) waitVolumeAttach(volumeID, instanceID string) error {
	return nil
}

func (d *driver) waitVolumeDetach(volumeID string) error {
	return nil
}

func (d *driver) RemoveVolume(volumeID string) error {
	return nil
}

func (d *driver) AttachVolume(
	runAsync bool,
	volumeID, instanceID string) ([]*core.VolumeAttachment, error) {
	log.WithField("provider", providerName).Debug("AttachVolume")
	return nil, nil

}

func (d *driver) DetachVolume(
	runAsync bool,
	volumeID, blank string) error {
	log.WithField("provider", providerName).Debug("DetachVolume")
	return nil
}

func (d *driver) CopySnapshot(runAsync bool,
	volumeID, snapshotID, snapshotName, destinationSnapshotName,
	destinationRegion string) (*core.Snapshot, error) {
	log.WithField("provider", providerName).Debug("CopySnapshot")
	return nil, nil
}

func configRegistration() *config.Registration {
	r := config.NewRegistration("Google GCE")
	r.Key(config.String, "", "", "", "gce.zone")
	r.Key(config.String, "", "", "", "gce.project")
	r.Key(config.String, "", "", "", "gce.keyfile")
	return r
}
<|endoftext|>


<commit_before>package turnpike

import (
	"fmt"
	"time"
)

const (
	defaultAuthTimeout = 2 * time.Minute
)

// A Realm is a WAMP routing and administrative domain.
//
// Clients that have connected to a WAMP router are joined to a realm and all
// message delivery is handled by the realm.
type Realm struct {
	_   string
	URI URI
	Broker
	Dealer
	Authorizer
	Interceptor
	CRAuthenticators map[string]CRAuthenticator
	Authenticators   map[string]Authenticator
	// DefaultAuth      func(details map[string]interface{}) (map[string]interface{}, error)
	AuthTimeout time.Duration
	clients     map[ID]Session
	localClient
	acts chan func()
}

type localClient struct {
	*Client
}

func (r *Realm) getPeer(details map[string]interface{}) (Peer, error) {
	peerA, peerB := localPipe()
	sess := Session{Peer: peerA, Id: NewID(), Details: details, kill: make(chan URI, 1)}
	if details == nil {
		details = make(map[string]interface{})
	}
	go r.handleSession(sess)
	log.Println("Established internal session:", sess)
	return peerB, nil
}

// Close disconnects all clients after sending a goodbye message
func (r Realm) Close() {
	r.acts <- func() {
		for _, client := range r.clients {
			client.kill <- ErrSystemShutdown
		}
	}

	var (
		sync     = make(chan struct{})
		nclients int
	)
	for {
		r.acts <- func() {
			nclients = len(r.clients)
			sync <- struct{}{}
		}
		<-sync
		if nclients == 0 {
			break
		}
	}

	close(r.acts)
}

func (r *Realm) init() {
	r.clients = make(map[ID]Session)
	r.acts = make(chan func())
	p, _ := r.getPeer(nil)
	r.localClient.Client = NewClient(p)
	if r.Broker == nil {
		r.Broker = NewDefaultBroker()
	}
	if r.Dealer == nil {
		r.Dealer = NewDefaultDealer()
	}
	if r.Authorizer == nil {
		r.Authorizer = NewDefaultAuthorizer()
	}
	if r.Interceptor == nil {
		r.Interceptor = NewDefaultInterceptor()
	}
	if r.AuthTimeout == 0 {
		r.AuthTimeout = defaultAuthTimeout
	}
	go r.run()
}

func (r *Realm) run() {
	for {
		if act, ok := <-r.acts; ok {
			act()
		} else {
			return
		}
	}
}

// func (r *Realm) metaHandler(c *Client) {
// }

func (l *localClient) onJoin(details map[string]interface{}) {
	l.Publish("wamp.session.on_join", []interface{}{details}, nil)
}

func (l *localClient) onLeave(session ID) {
	l.Publish("wamp.session.on_leave", []interface{}{session}, nil)
}

func (r *Realm) handleSession(sess Session) {
	sync := make(chan struct{})
	r.acts <- func() {
		r.clients[sess.Id] = sess
		r.onJoin(sess.Details)
		sync <- struct{}{}
	}
	<-sync
	defer func() {
		r.acts <- func() {
			delete(r.clients, sess.Id)
			r.Dealer.RemovePeer(sess.Peer)
			r.onLeave(sess.Id)
		}
	}()
	c := sess.Receive()
	// TODO: what happens if the realm is closed?

	for {
		var msg Message
		var open bool
		select {
		case msg, open = <-c:
			if !open {
				log.Println("lost session:", sess)
				return
			}
		case reason := <-sess.kill:
			logErr(sess.Send(&Goodbye{Reason: reason, Details: make(map[string]interface{})}))
			log.Printf("kill session %s: %v", sess, reason)
			// TODO: wait for client Goodbye?
			return
		}

		log.Printf("[%s] %s: %+v", sess, msg.MessageType(), msg)
		if isAuthz, err := r.Authorizer.Authorize(sess, msg); !isAuthz {
			errMsg := &Error{Type: msg.MessageType()}
			if err != nil {
				errMsg.Error = ErrAuthorizationFailed
				log.Printf("[%s] authorization failed: %v", sess, err)
			} else {
				errMsg.Error = ErrNotAuthorized
				log.Printf("[%s] %s UNAUTHORIZED", sess, msg.MessageType())
			}
			logErr(sess.Send(errMsg))
			continue
		}

		r.Interceptor.Intercept(sess, &msg)

		switch msg := msg.(type) {
		case *Goodbye:
			logErr(sess.Send(&Goodbye{Reason: ErrGoodbyeAndOut, Details: make(map[string]interface{})}))
			log.Printf("[%s] leaving: %v", sess, msg.Reason)
			return

		// Broker messages
		case *Publish:
			r.Broker.Publish(sess.Peer, msg)
		case *Subscribe:
			r.Broker.Subscribe(sess.Peer, msg)
		case *Unsubscribe:
			r.Broker.Unsubscribe(sess.Peer, msg)

		// Dealer messages
		case *Register:
			r.Dealer.Register(sess.Peer, msg)
		case *Unregister:
			r.Dealer.Unregister(sess.Peer, msg)
		case *Call:
			r.Dealer.Call(sess.Peer, msg)
		case *Yield:
			r.Dealer.Yield(sess.Peer, msg)

		// Error messages
		case *Error:
			if msg.Type == INVOCATION {
				// the only type of ERROR message the router should receive
				r.Dealer.Error(sess.Peer, msg)
			} else {
				log.Printf("invalid ERROR message received: %v", msg)
			}

		default:
			log.Println("Unhandled message:", msg.MessageType())
		}
	}
}

func (r *Realm) handleAuth(client Peer, details map[string]interface{}) (*Welcome, error) {
	msg, err := r.authenticate(details)
	if err != nil {
		return nil, err
	}
	// we should never get anything besides WELCOME and CHALLENGE
	if msg.MessageType() == WELCOME {
		return msg.(*Welcome), nil
	}
	// Challenge response
	challenge := msg.(*Challenge)
	if err := client.Send(challenge); err != nil {
		return nil, err
	}

	msg, err = GetMessageTimeout(client, r.AuthTimeout)
	if err != nil {
		return nil, err
	}
	log.Printf("%s: %+v", msg.MessageType(), msg)
	if authenticate, ok := msg.(*Authenticate); !ok {
		return nil, fmt.Errorf("unexpected %s message received", msg.MessageType())
	} else {
		return r.checkResponse(challenge, authenticate)
	}
}

// Authenticate either authenticates a client or returns a challenge message if
// challenge/response authentication is to be used.
func (r Realm) authenticate(details map[string]interface{}) (Message, error) {
	log.Println("details:", details)
	if len(r.Authenticators) == 0 && len(r.CRAuthenticators) == 0 {
		return &Welcome{}, nil
	}
	// TODO: this might not always be a []interface{}. Using the JSON unmarshaller it will be,
	// but we may have serializations that preserve more of the original type.
	// For now, the tests just explicitly send a []interface{}
	_authmethods, ok := details["authmethods"].([]interface{})
	if !ok {
		return nil, fmt.Errorf("No authentication supplied")
	}
	authmethods := []string{}
	for _, method := range _authmethods {
		if m, ok := method.(string); ok {
			authmethods = append(authmethods, m)
		} else {
			log.Printf("invalid authmethod value: %v", method)
		}
	}
	for _, method := range authmethods {
		if auth, ok := r.CRAuthenticators[method]; ok {
			if challenge, err := auth.Challenge(details); err != nil {
				return nil, err
			} else {
				return &Challenge{AuthMethod: method, Extra: challenge}, nil
			}
		}
		if auth, ok := r.Authenticators[method]; ok {
			if authDetails, err := auth.Authenticate(details); err != nil {
				return nil, err
			} else {
				return &Welcome{Details: addAuthMethod(authDetails, method)}, nil
			}
		}
	}
	// TODO: check default auth (special '*' auth?)
	return nil, fmt.Errorf("could not authenticate with any method")
}

// checkResponse determines whether the response to the challenge is sufficient to gain access to the Realm.
func (r Realm) checkResponse(chal *Challenge, auth *Authenticate) (*Welcome, error) {
	authenticator, ok := r.CRAuthenticators[chal.AuthMethod]
	if !ok {
		return nil, fmt.Errorf("authentication method has been removed")
	}
	if details, err := authenticator.Authenticate(chal.Extra, auth.Signature); err != nil {
		return nil, err
	} else {
		return &Welcome{Details: addAuthMethod(details, chal.AuthMethod)}, nil
	}
}

func addAuthMethod(details map[string]interface{}, method string) map[string]interface{} {
	if details == nil {
		details = make(map[string]interface{})
	}
	details["authmethod"] = method
	return details
}

// r := Realm{
// 	Authenticators: map[string]turnpike.Authenticator{
// 		"wampcra": turnpike.NewCRAAuthenticatorFactoryFactory(mySecret),
// 		"ticket": turnpike.NewTicketAuthenticator(myTicket),
// 		"asdfasdf": myAsdfAuthenticator,
// 	},
// 	BasicAuthenticators: map[string]turnpike.BasicAuthenticator{
// 		"anonymous": nil,
// 	},
// }
<commit_msg>Add Request ID to Error message during auth failures<commit_after>package turnpike

import (
	"fmt"
	"time"
)

const (
	defaultAuthTimeout = 2 * time.Minute
)

// A Realm is a WAMP routing and administrative domain.
//
// Clients that have connected to a WAMP router are joined to a realm and all
// message delivery is handled by the realm.
type Realm struct {
	_   string
	URI URI
	Broker
	Dealer
	Authorizer
	Interceptor
	CRAuthenticators map[string]CRAuthenticator
	Authenticators   map[string]Authenticator
	// DefaultAuth      func(details map[string]interface{}) (map[string]interface{}, error)
	AuthTimeout time.Duration
	clients     map[ID]Session
	localClient
	acts chan func()
}

type localClient struct {
	*Client
}

func (r *Realm) getPeer(details map[string]interface{}) (Peer, error) {
	peerA, peerB := localPipe()
	sess := Session{Peer: peerA, Id: NewID(), Details: details, kill: make(chan URI, 1)}
	if details == nil {
		details = make(map[string]interface{})
	}
	go r.handleSession(sess)
	log.Println("Established internal session:", sess)
	return peerB, nil
}

// Close disconnects all clients after sending a goodbye message
func (r Realm) Close() {
	r.acts <- func() {
		for _, client := range r.clients {
			client.kill <- ErrSystemShutdown
		}
	}

	var (
		sync     = make(chan struct{})
		nclients int
	)
	for {
		r.acts <- func() {
			nclients = len(r.clients)
			sync <- struct{}{}
		}
		<-sync
		if nclients == 0 {
			break
		}
	}

	close(r.acts)
}

func (r *Realm) init() {
	r.clients = make(map[ID]Session)
	r.acts = make(chan func())
	p, _ := r.getPeer(nil)
	r.localClient.Client = NewClient(p)
	if r.Broker == nil {
		r.Broker = NewDefaultBroker()
	}
	if r.Dealer == nil {
		r.Dealer = NewDefaultDealer()
	}
	if r.Authorizer == nil {
		r.Authorizer = NewDefaultAuthorizer()
	}
	if r.Interceptor == nil {
		r.Interceptor = NewDefaultInterceptor()
	}
	if r.AuthTimeout == 0 {
		r.AuthTimeout = defaultAuthTimeout
	}
	go r.run()
}

func (r *Realm) run() {
	for {
		if act, ok := <-r.acts; ok {
			act()
		} else {
			return
		}
	}
}

// func (r *Realm) metaHandler(c *Client) {
// }

func (l *localClient) onJoin(details map[string]interface{}) {
	l.Publish("wamp.session.on_join", []interface{}{details}, nil)
}

func (l *localClient) onLeave(session ID) {
	l.Publish("wamp.session.on_leave", []interface{}{session}, nil)
}

func (r *Realm) handleSession(sess Session) {
	sync := make(chan struct{})
	r.acts <- func() {
		r.clients[sess.Id] = sess
		r.onJoin(sess.Details)
		sync <- struct{}{}
	}
	<-sync
	defer func() {
		r.acts <- func() {
			delete(r.clients, sess.Id)
			r.Dealer.RemovePeer(sess.Peer)
			r.onLeave(sess.Id)
		}
	}()
	c := sess.Receive()
	// TODO: what happens if the realm is closed?

	for {
		var msg Message
		var open bool
		select {
		case msg, open = <-c:
			if !open {
				log.Println("lost session:", sess)
				return
			}
		case reason := <-sess.kill:
			logErr(sess.Send(&Goodbye{Reason: reason, Details: make(map[string]interface{})}))
			log.Printf("kill session %s: %v", sess, reason)
			// TODO: wait for client Goodbye?
			return
		}

		log.Printf("[%s] %s: %+v", sess, msg.MessageType(), msg)
		if isAuthz, err := r.Authorizer.Authorize(sess, msg); !isAuthz {
			errMsg := &Error{Type: msg.MessageType()}
			switch msg := msg.(type) {
			case *Publish:
				errMsg.Request = msg.Request
			case *Subscribe:
				errMsg.Request = msg.Request
			case *Unsubscribe:
				errMsg.Request = msg.Request
			case *Register:
				errMsg.Request = msg.Request
			case *Unregister:
				errMsg.Request = msg.Request
			case *Call:
				errMsg.Request = msg.Request
			case *Yield:
				errMsg.Request = msg.Request
			}
			if err != nil {
				errMsg.Error = ErrAuthorizationFailed
				log.Printf("[%s] authorization failed: %v", sess, err)
			} else {
				errMsg.Error = ErrNotAuthorized
				log.Printf("[%s] %s UNAUTHORIZED", sess, msg.MessageType())
			}
			logErr(sess.Send(errMsg))
			continue
		}

		r.Interceptor.Intercept(sess, &msg)

		switch msg := msg.(type) {
		case *Goodbye:
			logErr(sess.Send(&Goodbye{Reason: ErrGoodbyeAndOut, Details: make(map[string]interface{})}))
			log.Printf("[%s] leaving: %v", sess, msg.Reason)
			return

		// Broker messages
		case *Publish:
			r.Broker.Publish(sess.Peer, msg)
		case *Subscribe:
			r.Broker.Subscribe(sess.Peer, msg)
		case *Unsubscribe:
			r.Broker.Unsubscribe(sess.Peer, msg)

		// Dealer messages
		case *Register:
			r.Dealer.Register(sess.Peer, msg)
		case *Unregister:
			r.Dealer.Unregister(sess.Peer, msg)
		case *Call:
			r.Dealer.Call(sess.Peer, msg)
		case *Yield:
			r.Dealer.Yield(sess.Peer, msg)

		// Error messages
		case *Error:
			if msg.Type == INVOCATION {
				// the only type of ERROR message the router should receive
				r.Dealer.Error(sess.Peer, msg)
			} else {
				log.Printf("invalid ERROR message received: %v", msg)
			}

		default:
			log.Println("Unhandled message:", msg.MessageType())
		}
	}
}

func (r *Realm) handleAuth(client Peer, details map[string]interface{}) (*Welcome, error) {
	msg, err := r.authenticate(details)
	if err != nil {
		return nil, err
	}
	// we should never get anything besides WELCOME and CHALLENGE
	if msg.MessageType() == WELCOME {
		return msg.(*Welcome), nil
	}
	// Challenge response
	challenge := msg.(*Challenge)
	if err := client.Send(challenge); err != nil {
		return nil, err
	}

	msg, err = GetMessageTimeout(client, r.AuthTimeout)
	if err != nil {
		return nil, err
	}
	log.Printf("%s: %+v", msg.MessageType(), msg)
	if authenticate, ok := msg.(*Authenticate); !ok {
		return nil, fmt.Errorf("unexpected %s message received", msg.MessageType())
	} else {
		return r.checkResponse(challenge, authenticate)
	}
}

// Authenticate either authenticates a client or returns a challenge message if
// challenge/response authentication is to be used.
func (r Realm) authenticate(details map[string]interface{}) (Message, error) {
	log.Println("details:", details)
	if len(r.Authenticators) == 0 && len(r.CRAuthenticators) == 0 {
		return &Welcome{}, nil
	}
	// TODO: this might not always be a []interface{}. Using the JSON unmarshaller it will be,
	// but we may have serializations that preserve more of the original type.
	// For now, the tests just explicitly send a []interface{}
	_authmethods, ok := details["authmethods"].([]interface{})
	if !ok {
		return nil, fmt.Errorf("No authentication supplied")
	}
	authmethods := []string{}
	for _, method := range _authmethods {
		if m, ok := method.(string); ok {
			authmethods = append(authmethods, m)
		} else {
			log.Printf("invalid authmethod value: %v", method)
		}
	}
	for _, method := range authmethods {
		if auth, ok := r.CRAuthenticators[method]; ok {
			if challenge, err := auth.Challenge(details); err != nil {
				return nil, err
			} else {
				return &Challenge{AuthMethod: method, Extra: challenge}, nil
			}
		}
		if auth, ok := r.Authenticators[method]; ok {
			if authDetails, err := auth.Authenticate(details); err != nil {
				return nil, err
			} else {
				return &Welcome{Details: addAuthMethod(authDetails, method)}, nil
			}
		}
	}
	// TODO: check default auth (special '*' auth?)
	return nil, fmt.Errorf("could not authenticate with any method")
}

// checkResponse determines whether the response to the challenge is sufficient to gain access to the Realm.
func (r Realm) checkResponse(chal *Challenge, auth *Authenticate) (*Welcome, error) {
	authenticator, ok := r.CRAuthenticators[chal.AuthMethod]
	if !ok {
		return nil, fmt.Errorf("authentication method has been removed")
	}
	if details, err := authenticator.Authenticate(chal.Extra, auth.Signature); err != nil {
		return nil, err
	} else {
		return &Welcome{Details: addAuthMethod(details, chal.AuthMethod)}, nil
	}
}

func addAuthMethod(details map[string]interface{}, method string) map[string]interface{} {
	if details == nil {
		details = make(map[string]interface{})
	}
	details["authmethod"] = method
	return details
}

// r := Realm{
// 	Authenticators: map[string]turnpike.Authenticator{
// 		"wampcra": turnpike.NewCRAAuthenticatorFactoryFactory(mySecret),
// 		"ticket": turnpike.NewTicketAuthenticator(myTicket),
// 		"asdfasdf": myAsdfAuthenticator,
// 	},
// 	BasicAuthenticators: map[string]turnpike.BasicAuthenticator{
// 		"anonymous": nil,
// 	},
// }
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/howeyc/fsnotify"
	"go/build"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
)

func install(buildpath, lastError string) (installed bool, errorOutput string, err error) {
	cmdline := []string{"go", "install", "-v", buildpath}

	cmd := exec.Command("go", cmdline[1:]...)
	bufOut := bytes.NewBuffer([]byte{})
	bufErr := bytes.NewBuffer([]byte{})
	cmd.Stdout = bufOut
	cmd.Stderr = bufErr

	err = cmd.Run()

	if bufOut.Len() != 0 {
		errorOutput = bufOut.String()
		if errorOutput != lastError {
			fmt.Print(bufOut)
		}
		err = errors.New("compile error")
		return
	}

	installed = bufErr.Len() != 0

	return
}

func run(binName, binPath string, args []string) (runch chan bool) {
	runch = make(chan bool)
	go func() {
		cmdline := append([]string{binName}, args...)
		var proc *os.Process
		for _ = range runch {
			if proc != nil {
				proc.Kill()
			}
			cmd := exec.Command(binPath, args...)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			log.Print(cmdline)
			cmd.Start()
			proc = cmd.Process
		}
	}()
	return
}

func getWatcher(buildpath string) (watcher *fsnotify.Watcher, err error) {
	watcher, err = fsnotify.NewWatcher()
	addToWatcher(watcher, buildpath, map[string]bool{})
	return
}

func addToWatcher(watcher *fsnotify.Watcher, importpath string, watching map[string]bool) {
	pkg, err := build.Import(importpath, "", 0)
	if err != nil {
		return
	}
	if pkg.Goroot {
		return
	}
	watcher.Watch(pkg.Dir)
	watching[importpath] = true
	for _, imp := range pkg.Imports {
		if !watching[imp] {
			addToWatcher(watcher, imp, watching)
		}
	}
}

func rerun(buildpath string, args []string) (err error) {
	log.Printf("setting up rerun for %s %v", buildpath, args)

	pkg, err := build.Import(buildpath, "", 0)
	if err != nil {
		return
	}

	if pkg.Name != "main" {
		err = errors.New(fmt.Sprintf("expected package %q, got %q", "main", pkg.Name))
	}

	_, binName := path.Split(buildpath)
	binPath := filepath.Join(pkg.BinDir, binName)

	runch := run(binName, binPath, args)

	var errorOutput string
	_, errorOutput, ierr := install(buildpath, errorOutput)
	if ierr == nil {
		runch <- true
	}

	watcher, err := getWatcher(buildpath)
	if err != nil {
		log.Fatal(err)
	}

	for {
		we := <-watcher.Event
		var installed bool
		installed, errorOutput, _ = install(buildpath, errorOutput)
		if installed {
			log.Print(we.Name)
			runch <- true
			/* rescan */
			// watcher.Close()
			// watcher, err = getWatcher(buildpath)
			// if err != nil {
			// 	log.Fatal(err)
			// }
		}
	}
	return
}

func main() {
	log.SetPrefix("rerun ")
	if len(os.Args) < 2 {
		log.Fatal("Usage: rerun <import path> [arg]*")
	}
	err := rerun(os.Args[1], os.Args[2:])
	if err != nil {
		log.Print(err)
	}
}
<commit_msg>flush events before creating a new watcher<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/howeyc/fsnotify"
	"go/build"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
)

func install(buildpath, lastError string) (installed bool, errorOutput string, err error) {
	cmdline := []string{"go", "install", "-v", buildpath}

	cmd := exec.Command("go", cmdline[1:]...)
	bufOut := bytes.NewBuffer([]byte{})
	bufErr := bytes.NewBuffer([]byte{})
	cmd.Stdout = bufOut
	cmd.Stderr = bufErr

	err = cmd.Run()

	if bufOut.Len() != 0 {
		errorOutput = bufOut.String()
		if errorOutput != lastError {
			fmt.Print(bufOut)
		}
		err = errors.New("compile error")
		return
	}

	installed = bufErr.Len() != 0

	return
}

func run(binName, binPath string, args []string) (runch chan bool) {
	runch = make(chan bool)
	go func() {
		cmdline := append([]string{binName}, args...)
		var proc *os.Process
		for _ = range runch {
			if proc != nil {
				proc.Kill()
			}
			cmd := exec.Command(binPath, args...)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			log.Print(cmdline)
			cmd.Start()
			proc = cmd.Process
		}
	}()
	return
}

func getWatcher(buildpath string) (watcher *fsnotify.Watcher, err error) {
	watcher, err = fsnotify.NewWatcher()
	addToWatcher(watcher, buildpath, map[string]bool{})
	return
}

func addToWatcher(watcher *fsnotify.Watcher, importpath string, watching map[string]bool) {
	pkg, err := build.Import(importpath, "", 0)
	if err != nil {
		return
	}
	if pkg.Goroot {
		return
	}
	watcher.Watch(pkg.Dir)
	watching[importpath] = true
	for _, imp := range pkg.Imports {
		if !watching[imp] {
			addToWatcher(watcher, imp, watching)
		}
	}
}

func rerun(buildpath string, args []string) (err error) {
	log.Printf("setting up %s %v", buildpath, args)

	pkg, err := build.Import(buildpath, "", 0)
	if err != nil {
		return
	}

	if pkg.Name != "main" {
		err = errors.New(fmt.Sprintf("expected package %q, got %q", "main", pkg.Name))
		return
	}

	_, binName := path.Split(buildpath)
	binPath := filepath.Join(pkg.BinDir, binName)

	runch := run(binName, binPath, args)

	var errorOutput string
	_, errorOutput, ierr := install(buildpath, errorOutput)
	if ierr == nil {
		runch <- true
	}

	var watcher *fsnotify.Watcher
	watcher, err = getWatcher(buildpath)
	if err != nil {
		return
	}

	for {
		we := <-watcher.Event
		var installed bool
		installed, errorOutput, _ = install(buildpath, errorOutput)
		if installed {
			log.Print(we.Name)
			runch <- true
			watcher.Close()
			/* empty the buffer */
			for _ = range watcher.Event {

			}
			/* rescan */
			log.Println("rescanning")
			watcher, err = getWatcher(buildpath)
			if err != nil {
				return
			}
		}
	}
	return
}

func main() {
	if len(os.Args) < 2 {
		log.Fatal("Usage: rerun <import path> [arg]*")
	}
	err := rerun(os.Args[1], os.Args[2:])
	if err != nil {
		log.Print(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package testbase

import (
	"fmt"
	"time"

	gc "launchpad.net/gocheck"
	"launchpad.net/loggo"
)

// LoggingSuite redirects the juju logger to the test logger
// when embedded in a gocheck suite type.
type LoggingSuite struct {
	CleanupSuite
}

type gocheckWriter struct {
	c *gc.C
}

func (w *gocheckWriter) Write(level loggo.Level, module, filename string, line int, timestamp time.Time, message string) {
	// Magic calldepth value...
	w.c.Output(3, fmt.Sprintf("%s %s %s", level, module, message))
}

func (t *LoggingSuite) SetUpSuite(c *gc.C) {
	t.CleanupSuite.SetUpSuite(c)
	t.setUp(c)
	t.AddSuiteCleanup(func(*gc.C) {
		loggo.ResetLoggers()
		loggo.ResetWriters()
	})
}

func (t *LoggingSuite) SetUpTest(c *gc.C) {
	t.CleanupSuite.SetUpTest(c)
	t.PatchEnvironment("JUJU_LOGGING_CONFIG", "")
	t.setUp(c)
}

func (t *LoggingSuite) setUp(c *gc.C) {
	loggo.ResetWriters()
	loggo.ReplaceDefaultWriter(&gocheckWriter{c})
	loggo.ResetLoggers()
	loggo.GetLogger("juju").SetLogLevel(loggo.DEBUG)
}
<commit_msg>Also default the unit logging.<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package testbase

import (
	"fmt"
	"time"

	gc "launchpad.net/gocheck"
	"launchpad.net/loggo"
)

// LoggingSuite redirects the juju logger to the test logger
// when embedded in a gocheck suite type.
type LoggingSuite struct {
	CleanupSuite
}

type gocheckWriter struct {
	c *gc.C
}

func (w *gocheckWriter) Write(level loggo.Level, module, filename string, line int, timestamp time.Time, message string) {
	// Magic calldepth value...
	w.c.Output(3, fmt.Sprintf("%s %s %s", level, module, message))
}

func (t *LoggingSuite) SetUpSuite(c *gc.C) {
	t.CleanupSuite.SetUpSuite(c)
	t.setUp(c)
	t.AddSuiteCleanup(func(*gc.C) {
		loggo.ResetLoggers()
		loggo.ResetWriters()
	})
}

func (t *LoggingSuite) SetUpTest(c *gc.C) {
	t.CleanupSuite.SetUpTest(c)
	t.PatchEnvironment("JUJU_LOGGING_CONFIG", "")
	t.setUp(c)
}

func (t *LoggingSuite) setUp(c *gc.C) {
	loggo.ResetWriters()
	loggo.ReplaceDefaultWriter(&gocheckWriter{c})
	loggo.ResetLoggers()
	loggo.GetLogger("juju").SetLogLevel(loggo.DEBUG)
	loggo.GetLogger("unit").SetLogLevel(loggo.DEBUG)
}
<|endoftext|>


<commit_before>// The RS232 package lets you access old serial junk from go
package rs232

/*
#include <stdlib.h>
#include <fcntl.h>
#include <termios.h>

void initBaudRates();
*/
import "C"

import (
	"os"
	"syscall"
)

var baudConversionMap = map[int]_Ctype_speed_t{}

// This is your serial port handle.
type SerialPort struct {
	port *os.File
}

func init() {
	C.initBaudRates()
}

//export addBaudRate
func addBaudRate(num int, val _Ctype_speed_t) {
	baudConversionMap[num] = val
}

func baudConversion(rate int) (flag _Ctype_speed_t) {
	return baudConversionMap[rate]
}

// SerConf represents the basic serial configuration to provide to OpenPort.
type SerConf int

const (
	S_8N1 = iota
	S_7E1
	S_7O1
)

// Opens and returns a non-blocking serial port.
// The device, baud rate, and SerConf is specified.
//
// Example:  rs232.OpenPort("/dev/ttyS0", 115200, rs232.S_8N1)
func OpenPort(port string, baudRate int, serconf SerConf) (rv *SerialPort, err error) {
	rv = &SerialPort{}
	f, open_err := os.OpenFile(port,
		syscall.O_RDWR|syscall.O_NOCTTY|syscall.O_NDELAY,
		0666)
	if open_err != nil {
		err = open_err
		return
	}
	rv.port = f

	fd := rv.port.Fd()

	var options C.struct_termios
	if C.tcgetattr(C.int(fd), &options) < 0 {
		panic("tcgetattr failed")
	}

	if C.cfsetispeed(&options, baudConversion(baudRate)) < 0 {
		panic("cfsetispeed failed")
	}
	if C.cfsetospeed(&options, baudConversion(baudRate)) < 0 {
		panic("cfsetospeed failed")
	}
	switch serconf {
	case S_8N1:
		{
			options.c_cflag &^= C.PARENB
			options.c_cflag &^= C.CSTOPB
			options.c_cflag &^= C.CSIZE
			options.c_cflag |= C.CS8
		}
	case S_7E1:
		{
			options.c_cflag |= C.PARENB
			options.c_cflag &^= C.PARODD
			options.c_cflag &^= C.CSTOPB
			options.c_cflag &^= C.CSIZE
			options.c_cflag |= C.CS7
		}
	case S_7O1:
		{
			options.c_cflag |= C.PARENB
			options.c_cflag |= C.PARODD
			options.c_cflag &^= C.CSTOPB
			options.c_cflag &^= C.CSIZE
			options.c_cflag |= C.CS7
		}
	}
	// Local
	options.c_cflag |= (C.CLOCAL | C.CREAD)
	// no hardware flow control
	options.c_cflag &^= C.CRTSCTS
	// Don't EOF on a zero read, just block
	options.c_cc[C.VMIN] = 1

	if C.tcsetattr(C.int(fd), C.TCSANOW, &options) < 0 {
		panic("tcsetattr failed")
	}

	if syscall.SetNonblock(int(fd), false) != nil {
		panic("Error disabling blocking")
	}

	return
}

// Read from the port.
func (port *SerialPort) Read(p []byte) (n int, err error) {
	return port.port.Read(p)
}

// Write to the port.
func (port *SerialPort) Write(p []byte) (n int, err error) {
	return port.port.Write(p)
}

// Close the port.
func (port *SerialPort) Close() error {
	return port.port.Close()
}
<commit_msg>Make the SerConf constants be SerConfs.<commit_after>// The RS232 package lets you access old serial junk from go
package rs232

/*
#include <stdlib.h>
#include <fcntl.h>
#include <termios.h>

void initBaudRates();
*/
import "C"

import (
	"os"
	"syscall"
)

var baudConversionMap = map[int]_Ctype_speed_t{}

// This is your serial port handle.
type SerialPort struct {
	port *os.File
}

func init() {
	C.initBaudRates()
}

//export addBaudRate
func addBaudRate(num int, val _Ctype_speed_t) {
	baudConversionMap[num] = val
}

func baudConversion(rate int) (flag _Ctype_speed_t) {
	return baudConversionMap[rate]
}

// SerConf represents the basic serial configuration to provide to OpenPort.
type SerConf int

const (
	S_8N1 SerConf = iota
	S_7E1
	S_7O1
)

// Opens and returns a non-blocking serial port.
// The device, baud rate, and SerConf is specified.
//
// Example:  rs232.OpenPort("/dev/ttyS0", 115200, rs232.S_8N1)
func OpenPort(port string, baudRate int, serconf SerConf) (rv *SerialPort, err error) {
	rv = &SerialPort{}
	f, open_err := os.OpenFile(port,
		syscall.O_RDWR|syscall.O_NOCTTY|syscall.O_NDELAY,
		0666)
	if open_err != nil {
		err = open_err
		return
	}
	rv.port = f

	fd := rv.port.Fd()

	var options C.struct_termios
	if C.tcgetattr(C.int(fd), &options) < 0 {
		panic("tcgetattr failed")
	}

	if C.cfsetispeed(&options, baudConversion(baudRate)) < 0 {
		panic("cfsetispeed failed")
	}
	if C.cfsetospeed(&options, baudConversion(baudRate)) < 0 {
		panic("cfsetospeed failed")
	}
	switch serconf {
	case S_8N1:
		{
			options.c_cflag &^= C.PARENB
			options.c_cflag &^= C.CSTOPB
			options.c_cflag &^= C.CSIZE
			options.c_cflag |= C.CS8
		}
	case S_7E1:
		{
			options.c_cflag |= C.PARENB
			options.c_cflag &^= C.PARODD
			options.c_cflag &^= C.CSTOPB
			options.c_cflag &^= C.CSIZE
			options.c_cflag |= C.CS7
		}
	case S_7O1:
		{
			options.c_cflag |= C.PARENB
			options.c_cflag |= C.PARODD
			options.c_cflag &^= C.CSTOPB
			options.c_cflag &^= C.CSIZE
			options.c_cflag |= C.CS7
		}
	}
	// Local
	options.c_cflag |= (C.CLOCAL | C.CREAD)
	// no hardware flow control
	options.c_cflag &^= C.CRTSCTS
	// Don't EOF on a zero read, just block
	options.c_cc[C.VMIN] = 1

	if C.tcsetattr(C.int(fd), C.TCSANOW, &options) < 0 {
		panic("tcsetattr failed")
	}

	if syscall.SetNonblock(int(fd), false) != nil {
		panic("Error disabling blocking")
	}

	return
}

// Read from the port.
func (port *SerialPort) Read(p []byte) (n int, err error) {
	return port.port.Read(p)
}

// Write to the port.
func (port *SerialPort) Write(p []byte) (n int, err error) {
	return port.port.Write(p)
}

// Close the port.
func (port *SerialPort) Close() error {
	return port.port.Close()
}
<|endoftext|>


<commit_before>// +build watch

// Package watch provides a filesystem watcher that is used to rebuild affected targets.
package watch

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/kardianos/osext"
	"github.com/streamrail/concurrent-map"
	"gopkg.in/op/go-logging.v1"

	"core"
)

var log = logging.MustGetLogger("watch")

const debounceInterval = 50 * time.Millisecond

// Watch starts watching the sources of the given labels for changes and triggers
// rebuilds whenever they change.
// It never returns successfully, it will either watch forever or die.
func Watch(state *core.BuildState, labels []core.BuildLabel) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatalf("Error setting up watcher: %s", err)
	}
	// This sets up the actual watches. It must be done in a separate goroutine.
	files := cmap.New()
	go startWatching(watcher, state, labels, files)

	// If any of the targets are tests, we'll run plz test, otherwise just plz build.
	command := "build"
	for _, label := range labels {
		if state.Graph.TargetOrDie(label).IsTest {
			command = "test"
			break
		}
	}
	log.Notice("Command: %s", command)

	for {
		select {
		case event := <-watcher.Events:
			log.Info("Event: %s", event)
			if !files.Has(event.Name) {
				log.Notice("Skipping notification for %s", event.Name)
				continue
			}
			// Quick debounce; poll and discard all events for the next brief period.
		outer:
			for {
				select {
				case <-watcher.Events:
				case <-time.After(debounceInterval):
					break outer
				}
			}
			runBuild(state, command, labels)
		case err := <-watcher.Errors:
			log.Error("Error watching files:", err)
		}
	}
}

func runBuild(state *core.BuildState, command string, labels []core.BuildLabel) {
	binary, err := osext.Executable()
	if err != nil {
		log.Warning("Can't determine current executable, will assume 'plz'")
		binary = "plz"
	}
	cmd := exec.Command(binary, command)
	cmd.Args = append(cmd.Args, "-c", state.Config.Build.Config)
	for _, label := range labels {
		cmd.Args = append(cmd.Args, label.String())
	}
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	log.Notice("Running %s %s...", binary, command)
	if err := cmd.Run(); err != nil {
		// Only log the error if it's not a straightforward non-zero exit; the user will presumably
		// already have been pestered about that.
		if _, ok := err.(*exec.ExitError); !ok {
			log.Error("Failed to run %s: %s", binary, err)
		}
	}
}

func startWatching(watcher *fsnotify.Watcher, state *core.BuildState, labels []core.BuildLabel, files cmap.ConcurrentMap) {
	// Deduplicate seen targets & sources.
	targets := map[*core.BuildTarget]struct{}{}
	dirs := map[string]struct{}{}

	var startWatch func(*core.BuildTarget)
	startWatch = func(target *core.BuildTarget) {
		if _, present := targets[target]; present {
			return
		}
		targets[target] = struct{}{}
		for _, source := range target.AllSources() {
			if source.Label() == nil {
				for _, src := range source.Paths(state.Graph) {
					files.Set(src, struct{}{})
					if info, err := os.Stat(src); err == nil && !info.IsDir() {
						src = path.Dir(src)
					}
					if _, present := dirs[src]; !present {
						log.Notice("Adding watch on %s", src)
						dirs[src] = struct{}{}
						if err := watcher.Add(src); err != nil {
							log.Error("Failed to add watch on %s: %s", src, err)
						}
					}
				}
			}
		}
		for _, dep := range target.Dependencies() {
			startWatch(dep)
		}
		pkg := state.Graph.PackageOrDie(target.Label.PackageName)
		if !files.Has(pkg.Filename) {
			log.Notice("Adding watch on %s", pkg.Filename)
			files.Set(pkg.Filename, struct{}{})
		}
		for _, subinclude := range pkg.Subincludes {
			startWatch(state.Graph.TargetOrDie(subinclude))
		}
	}

	for _, label := range labels {
		startWatch(state.Graph.TargetOrDie(label))
	}
	// Drop a message here so they know when it's actually ready to go.
	fmt.Println("And now my watch begins...")
}
<commit_msg>Fix watching for data files and for srcs or data that are directories.<commit_after>// +build watch

// Package watch provides a filesystem watcher that is used to rebuild affected targets.
package watch

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/kardianos/osext"
	"github.com/streamrail/concurrent-map"
	"gopkg.in/op/go-logging.v1"

	"core"
)

var log = logging.MustGetLogger("watch")

const debounceInterval = 50 * time.Millisecond

// Watch starts watching the sources of the given labels for changes and triggers
// rebuilds whenever they change.
// It never returns successfully, it will either watch forever or die.
func Watch(state *core.BuildState, labels []core.BuildLabel) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatalf("Error setting up watcher: %s", err)
	}
	// This sets up the actual watches. It must be done in a separate goroutine.
	files := cmap.New()
	go startWatching(watcher, state, labels, files)

	// If any of the targets are tests, we'll run plz test, otherwise just plz build.
	command := "build"
	for _, label := range labels {
		if state.Graph.TargetOrDie(label).IsTest {
			command = "test"
			break
		}
	}
	log.Notice("Command: %s", command)

	for {
		select {
		case event := <-watcher.Events:
			log.Info("Event: %s", event)
			if !files.Has(event.Name) {
				log.Notice("Skipping notification for %s", event.Name)
				continue
			}
			// Quick debounce; poll and discard all events for the next brief period.
		outer:
			for {
				select {
				case <-watcher.Events:
				case <-time.After(debounceInterval):
					break outer
				}
			}
			runBuild(state, command, labels)
		case err := <-watcher.Errors:
			log.Error("Error watching files:", err)
		}
	}
}

func runBuild(state *core.BuildState, command string, labels []core.BuildLabel) {
	binary, err := osext.Executable()
	if err != nil {
		log.Warning("Can't determine current executable, will assume 'plz'")
		binary = "plz"
	}
	cmd := exec.Command(binary, command)
	cmd.Args = append(cmd.Args, "-c", state.Config.Build.Config)
	for _, label := range labels {
		cmd.Args = append(cmd.Args, label.String())
	}
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	log.Notice("Running %s %s...", binary, command)
	if err := cmd.Run(); err != nil {
		// Only log the error if it's not a straightforward non-zero exit; the user will presumably
		// already have been pestered about that.
		if _, ok := err.(*exec.ExitError); !ok {
			log.Error("Failed to run %s: %s", binary, err)
		}
	}
}

func startWatching(watcher *fsnotify.Watcher, state *core.BuildState, labels []core.BuildLabel, files cmap.ConcurrentMap) {
	// Deduplicate seen targets & sources.
	targets := map[*core.BuildTarget]struct{}{}
	dirs := map[string]struct{}{}

	var startWatch func(*core.BuildTarget)
	startWatch = func(target *core.BuildTarget) {
		if _, present := targets[target]; present {
			return
		}
		targets[target] = struct{}{}
		for _, source := range target.AllSources() {
			addSource(watcher, state, source, dirs, files)
		}
		for _, datum := range target.Data {
			addSource(watcher, state, datum, dirs, files)
		}
		for _, dep := range target.Dependencies() {
			startWatch(dep)
		}
		pkg := state.Graph.PackageOrDie(target.Label.PackageName)
		if !files.Has(pkg.Filename) {
			log.Notice("Adding watch on %s", pkg.Filename)
			files.Set(pkg.Filename, struct{}{})
		}
		for _, subinclude := range pkg.Subincludes {
			startWatch(state.Graph.TargetOrDie(subinclude))
		}
	}

	for _, label := range labels {
		startWatch(state.Graph.TargetOrDie(label))
	}
	// Drop a message here so they know when it's actually ready to go.
	fmt.Println("And now my watch begins...")
}

func addSource(watcher *fsnotify.Watcher, state *core.BuildState, source core.BuildInput, dirs map[string]struct{}, files cmap.ConcurrentMap) {
	if source.Label() == nil {
		for _, src := range source.Paths(state.Graph) {
			if err := filepath.Walk(src, func(src string, info os.FileInfo, err error) error {
				files.Set(src, struct{}{})
				dir := src
				if info, err := os.Stat(src); err == nil && !info.IsDir() {
					dir = path.Dir(src)
				}
				if _, present := dirs[dir]; !present {
					log.Notice("Adding watch on %s", dir)
					dirs[dir] = struct{}{}
					if err := watcher.Add(dir); err != nil {
						log.Error("Failed to add watch on %s: %s", src, err)
					}
				}
				return err
			}); err != nil {
				log.Error("Failed to add watch on %s: %s", src, err)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"

	"github.com/fatih/color"
	"gopkg.in/alecthomas/kingpin.v2"

	"github.com/InnovaCo/serve/manifest"
	_ "github.com/InnovaCo/serve/plugins"
)

func main() {
	manifestFile := kingpin.Flag("manifest", "Path to manifest.yml file.").Default("manifest.yml").String()
	plugin       := kingpin.Arg("plugin", "Plugin name for run.").String()
	vars         := *kingpin.Flag("var", "key=value pairs with manifest vars.").StringMap()
	dryRun 	 	 := kingpin.Flag("dry-run", "Show manifest section only").Bool()

	kingpin.Parse()

	mnf := manifest.Load(*manifestFile, vars)

	if (*plugin == "") && *dryRun {
		fmt.Println(mnf)
		return
	}

	plugins, err := mnf.FindPlugins(*plugin)
	if err != nil {
		log.Fatalf("Error find plugins for '%s': %v", *plugin, err)
	}

	for _, pair := range plugins {
		log.Println(color.GreenString("%v:\n", pair.PluginName), pair.Data)

		if !*dryRun {
			if err := pair.Plugin.Run(pair.Data); err != nil {
				fmt.Println("")
				log.Fatalln(color.RedString("Error on run plugin `%s`: %v", pair.PluginName, err))
			}
		}
	}
}
<commit_msg>= enable color output for all<commit_after>package main

import (
	"fmt"
	"log"

	"github.com/fatih/color"
	"gopkg.in/alecthomas/kingpin.v2"

	"github.com/InnovaCo/serve/manifest"
	_ "github.com/InnovaCo/serve/plugins"
)

func init() {
	color.NoColor = false
}

func main() {
	manifestFile := kingpin.Flag("manifest", "Path to manifest.yml file.").Default("manifest.yml").String()
	plugin       := kingpin.Arg("plugin", "Plugin name for run.").String()
	vars         := *kingpin.Flag("var", "key=value pairs with manifest vars.").StringMap()
	dryRun 	 	 := kingpin.Flag("dry-run", "Show manifest section only").Bool()

	kingpin.Parse()

	mnf := manifest.Load(*manifestFile, vars)

	if (*plugin == "") && *dryRun {
		fmt.Println(mnf)
		return
	}

	plugins, err := mnf.FindPlugins(*plugin)
	if err != nil {
		log.Fatalf("Error find plugins for '%s': %v", *plugin, err)
	}

	for _, pair := range plugins {
		log.Println(color.GreenString("%v:\n", pair.PluginName), pair.Data)

		if !*dryRun {
			if err := pair.Plugin.Run(pair.Data); err != nil {
				fmt.Println("")
				log.Fatalln(color.RedString("Error on run plugin `%s`: %v", pair.PluginName, err))
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"github.com/abh/geodns/countries"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
)

func getQuestionName(z *Zone, req *dns.Msg) string {
	lx := dns.SplitLabels(req.Question[0].Name)
	ql := lx[0 : len(lx)-z.LenLabels]
	return strings.ToLower(strings.Join(ql, "."))
}

var geoIP = setupGeoIP()

func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) {

	qtype := req.Question[0].Qtype

	logPrintf("[zone %s] incoming %s %s %d from %s\n", z.Origin, req.Question[0].Name,
		dns.TypeToString[qtype], req.MsgHdr.Id, w.RemoteAddr())

	qCounter.Add(1)
	logPrintln("Got request", req)

	label := getQuestionName(z, req)

	var ip string
	var edns *dns.EDNS0_SUBNET
	var opt_rr *dns.OPT

	for _, extra := range req.Extra {
		log.Println("Extra", extra)
		for _, o := range extra.(*dns.OPT).Option {
			opt_rr = extra.(*dns.OPT)
			switch e := o.(type) {
			case *dns.EDNS0_NSID:
				// do stuff with e.Nsid
			case *dns.EDNS0_SUBNET:
				log.Println("========== XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
				log.Println("Got edns", e.Address, e.Family, e.SourceNetmask, e.SourceScope)
				if e.Address != nil {
					log.Println("Setting edns to", e)
					edns = e
					ip = e.Address.String()
				}
			}
		}
	}

	var country string
	if geoIP != nil {
		if len(ip) == 0 { // no edns subnet
			ip, _, _ = net.SplitHostPort(w.RemoteAddr().String())
		}
		country = strings.ToLower(geoIP.GetCountry(ip))
		logPrintln("Country:", ip, country)
	}

	m := new(dns.Msg)
	m.SetReply(req)
	if e := m.IsEdns0(); e != nil {
		m.SetEdns0(4096, e.Do())
	}
	m.Authoritative = true

	// TODO: set scope to 0 if there are no alternate responses
	if edns != nil {
		log.Println("family", edns.Family)
		if edns.Family != 0 {
			log.Println("edns response!")
			edns.SourceScope = 16
			m.Extra = append(m.Extra, opt_rr)
		}
	}

	// TODO(ask) Fix the findLabels API to make this work better
	if alias := z.findLabels(label, "", dns.TypeMF); alias != nil &&
		alias.Records[dns.TypeMF] != nil {
		// We found an alias record, so pretend the question was for that name instead
		label = alias.firstRR(dns.TypeMF).(*dns.MF).Mf
	}

	labels := z.findLabels(label, country, qtype)
	if labels == nil {

		if label == "_status" && (qtype == dns.TypeANY || qtype == dns.TypeTXT) {
			m.Answer = statusRR(z)
			m.Authoritative = true
			w.WriteMsg(m)
			return
		}

		if label == "_country" && (qtype == dns.TypeANY || qtype == dns.TypeTXT) {
			h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
			h.Name = "_country." + z.Origin + "."

			m.Answer = []dns.RR{&dns.TXT{Hdr: h,
				Txt: []string{
					w.RemoteAddr().String(),
					ip,
					string(country),
					string(countries.CountryContinent[country]),
				},
			}}

			m.Authoritative = true
			w.WriteMsg(m)
			return
		}

		// return NXDOMAIN
		m.SetRcode(req, dns.RcodeNameError)
		m.Authoritative = true

		m.Ns = []dns.RR{z.SoaRR()}

		w.WriteMsg(m)
		return
	}

	if servers := labels.Picker(qtype, labels.MaxHosts); servers != nil {
		var rrs []dns.RR
		for _, record := range servers {
			rr := record.RR
			rr.Header().Name = req.Question[0].Name
			rrs = append(rrs, rr)
		}
		m.Answer = rrs
	}

	if len(m.Answer) == 0 {
		if labels := z.Labels[label]; labels != nil {
			if _, ok := labels.Records[dns.TypeCNAME]; ok {
				cname := labels.firstRR(dns.TypeCNAME)
				m.Answer = append(m.Answer, cname)
			} else {
				m.Ns = append(m.Ns, z.SoaRR())
			}
		} else {
			m.Ns = append(m.Ns, z.SoaRR())
		}
	}

	logPrintln(m)

	err := w.WriteMsg(m)
	if err != nil {
		// if Pack'ing fails the Write fails. Return SERVFAIL.
		log.Println("Error writing packet", m)
		dns.HandleFailed(w, req)
	}
	return
}

func statusRR(z *Zone) []dns.RR {
	h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
	h.Name = "_status." + z.Origin + "."

	status := map[string]string{"v": VERSION, "id": serverId}

	hostname, err := os.Hostname()
	if err == nil {
		status["h"] = hostname
	}
	status["up"] = strconv.Itoa(int(time.Since(timeStarted).Seconds()))
	status["qs"] = qCounter.String()

	js, err := json.Marshal(status)

	return []dns.RR{&dns.TXT{Hdr: h, Txt: []string{string(js)}}}
}

func setupServerFunc(Zone *Zone) func(dns.ResponseWriter, *dns.Msg) {
	return func(w dns.ResponseWriter, r *dns.Msg) {
		serve(w, r, Zone)
	}
}

func listenAndServe(ip string, Zones *Zones) {

	prots := []string{"udp", "tcp"}

	for _, prot := range prots {
		go func(p string) {
			server := &dns.Server{Addr: ip, Net: p}

			log.Printf("Opening on %s %s", ip, p)
			if err := server.ListenAndServe(); err != nil {
				log.Fatalf("geodns: failed to setup %s %s: %s", ip, p, err)
			}
			log.Fatalf("geodns: ListenAndServe unexpectedly returned")
		}(prot)
	}

}
<commit_msg>Remove a little debug noise<commit_after>package main

import (
	"encoding/json"
	"github.com/abh/geodns/countries"
	"github.com/miekg/dns"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
)

func getQuestionName(z *Zone, req *dns.Msg) string {
	lx := dns.SplitLabels(req.Question[0].Name)
	ql := lx[0 : len(lx)-z.LenLabels]
	return strings.ToLower(strings.Join(ql, "."))
}

var geoIP = setupGeoIP()

func serve(w dns.ResponseWriter, req *dns.Msg, z *Zone) {

	qtype := req.Question[0].Qtype

	logPrintf("[zone %s] incoming %s %s %d from %s\n", z.Origin, req.Question[0].Name,
		dns.TypeToString[qtype], req.MsgHdr.Id, w.RemoteAddr())

	qCounter.Add(1)

	logPrintln("Got request", req)

	label := getQuestionName(z, req)

	var ip string
	var edns *dns.EDNS0_SUBNET
	var opt_rr *dns.OPT

	for _, extra := range req.Extra {
		log.Println("Extra", extra)
		for _, o := range extra.(*dns.OPT).Option {
			opt_rr = extra.(*dns.OPT)
			switch e := o.(type) {
			case *dns.EDNS0_NSID:
				// do stuff with e.Nsid
			case *dns.EDNS0_SUBNET:
				log.Println("Got edns", e.Address, e.Family, e.SourceNetmask, e.SourceScope)
				if e.Address != nil {
					log.Println("Setting edns to", e)
					edns = e
					ip = e.Address.String()
				}
			}
		}
	}

	var country string
	if geoIP != nil {
		if len(ip) == 0 { // no edns subnet
			ip, _, _ = net.SplitHostPort(w.RemoteAddr().String())
		}
		country = strings.ToLower(geoIP.GetCountry(ip))
		logPrintln("Country:", ip, country)
	}

	m := new(dns.Msg)
	m.SetReply(req)
	if e := m.IsEdns0(); e != nil {
		m.SetEdns0(4096, e.Do())
	}
	m.Authoritative = true

	// TODO: set scope to 0 if there are no alternate responses
	if edns != nil {
		log.Println("family", edns.Family)
		if edns.Family != 0 {
			log.Println("edns response!")
			edns.SourceScope = 16
			m.Extra = append(m.Extra, opt_rr)
		}
	}

	// TODO(ask) Fix the findLabels API to make this work better
	if alias := z.findLabels(label, "", dns.TypeMF); alias != nil &&
		alias.Records[dns.TypeMF] != nil {
		// We found an alias record, so pretend the question was for that name instead
		label = alias.firstRR(dns.TypeMF).(*dns.MF).Mf
	}

	labels := z.findLabels(label, country, qtype)
	if labels == nil {

		if label == "_status" && (qtype == dns.TypeANY || qtype == dns.TypeTXT) {
			m.Answer = statusRR(z)
			m.Authoritative = true
			w.WriteMsg(m)
			return
		}

		if label == "_country" && (qtype == dns.TypeANY || qtype == dns.TypeTXT) {
			h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
			h.Name = "_country." + z.Origin + "."

			m.Answer = []dns.RR{&dns.TXT{Hdr: h,
				Txt: []string{
					w.RemoteAddr().String(),
					ip,
					string(country),
					string(countries.CountryContinent[country]),
				},
			}}

			m.Authoritative = true
			w.WriteMsg(m)
			return
		}

		// return NXDOMAIN
		m.SetRcode(req, dns.RcodeNameError)
		m.Authoritative = true

		m.Ns = []dns.RR{z.SoaRR()}

		w.WriteMsg(m)
		return
	}

	if servers := labels.Picker(qtype, labels.MaxHosts); servers != nil {
		var rrs []dns.RR
		for _, record := range servers {
			rr := record.RR
			rr.Header().Name = req.Question[0].Name
			rrs = append(rrs, rr)
		}
		m.Answer = rrs
	}

	if len(m.Answer) == 0 {
		if labels := z.Labels[label]; labels != nil {
			if _, ok := labels.Records[dns.TypeCNAME]; ok {
				cname := labels.firstRR(dns.TypeCNAME)
				m.Answer = append(m.Answer, cname)
			} else {
				m.Ns = append(m.Ns, z.SoaRR())
			}
		} else {
			m.Ns = append(m.Ns, z.SoaRR())
		}
	}

	logPrintln(m)

	err := w.WriteMsg(m)
	if err != nil {
		// if Pack'ing fails the Write fails. Return SERVFAIL.
		log.Println("Error writing packet", m)
		dns.HandleFailed(w, req)
	}
	return
}

func statusRR(z *Zone) []dns.RR {
	h := dns.RR_Header{Ttl: 1, Class: dns.ClassINET, Rrtype: dns.TypeTXT}
	h.Name = "_status." + z.Origin + "."

	status := map[string]string{"v": VERSION, "id": serverId}

	hostname, err := os.Hostname()
	if err == nil {
		status["h"] = hostname
	}
	status["up"] = strconv.Itoa(int(time.Since(timeStarted).Seconds()))
	status["qs"] = qCounter.String()

	js, err := json.Marshal(status)

	return []dns.RR{&dns.TXT{Hdr: h, Txt: []string{string(js)}}}
}

func setupServerFunc(Zone *Zone) func(dns.ResponseWriter, *dns.Msg) {
	return func(w dns.ResponseWriter, r *dns.Msg) {
		serve(w, r, Zone)
	}
}

func listenAndServe(ip string, Zones *Zones) {

	prots := []string{"udp", "tcp"}

	for _, prot := range prots {
		go func(p string) {
			server := &dns.Server{Addr: ip, Net: p}

			log.Printf("Opening on %s %s", ip, p)
			if err := server.ListenAndServe(); err != nil {
				log.Fatalf("geodns: failed to setup %s %s: %s", ip, p, err)
			}
			log.Fatalf("geodns: ListenAndServe unexpectedly returned")
		}(prot)
	}

}
<|endoftext|>


<commit_before>package mpower

const (
	baseLive = "https://app.mpowerpayments.com/api/v1"
	baseTest = "https://app.mpowerpayments.com/sandbox-api/v1"
)

// Setup as defined by mpower docs with the exception of the BaseURL
type Setup struct {
	MasterKey   string
	PrivateKey  string
	PublicKey   string
	Token       string
	ContentType string
	Headers     map[string]string
	BaseURL     string
}

// NewSetup - returns a new setup object
func NewSetup(masterKey, privateKey, publicKey, token string) *Setup {
	setup := &Setup{
		MasterKey:  masterKey,
		PrivateKey: privateKey,
		PublicKey:  publicKey,
		Token:      token,
	}

	setup.setupHeaders()

	return setup
}

// NewSetupFromEnv creates a setup from your environment keys
func NewSetupFromEnv() *Setup {
	setup := &Setup{
		MasterKey:  env("MP-Master-Key"),
		PrivateKey: env("MP-Private-Key"),
		PublicKey:  env("MP-Public-Key"),
		Token:      env("MP-Token"),
	}

	setup.setupHeaders()

	return setup
}

func (s *Setup) setupHeaders() {
	s.Headers = make(map[string]string)
	s.Headers["MP-Master-Key"] = s.MasterKey
	s.Headers["MP-Private-Key"] = s.PrivateKey
	s.Headers["MP-Public-Key"] = s.PublicKey
	s.Headers["MP-Token"] = s.Token
	s.Headers["Content-Type"] = "application/json"
}
<commit_msg>change env keys from hyphen to underscore<commit_after>package mpower

const (
	baseLive = "https://app.mpowerpayments.com/api/v1"
	baseTest = "https://app.mpowerpayments.com/sandbox-api/v1"
)

// Setup as defined by mpower docs with the exception of the BaseURL
type Setup struct {
	MasterKey   string
	PrivateKey  string
	PublicKey   string
	Token       string
	ContentType string
	Headers     map[string]string
	BaseURL     string
}

// NewSetup - returns a new setup object
func NewSetup(masterKey, privateKey, publicKey, token string) *Setup {
	setup := &Setup{
		MasterKey:  masterKey,
		PrivateKey: privateKey,
		PublicKey:  publicKey,
		Token:      token,
	}

	setup.setupHeaders()

	return setup
}

// NewSetupFromEnv creates a setup from your environment keys
func NewSetupFromEnv() *Setup {
	setup := &Setup{
		MasterKey:  env("MP_Master_Key"),
		PrivateKey: env("MP_Private_Key"),
		PublicKey:  env("MP_Public_Key"),
		Token:      env("MP_Token"),
	}

	setup.setupHeaders()

	return setup
}

func (s *Setup) setupHeaders() {
	s.Headers = make(map[string]string)
	s.Headers["MP-Master-Key"] = s.MasterKey
	s.Headers["MP-Private-Key"] = s.PrivateKey
	s.Headers["MP-Public-Key"] = s.PublicKey
	s.Headers["MP-Token"] = s.Token
	s.Headers["Content-Type"] = "application/json"
}
<|endoftext|>


<commit_before>package fest

import (
	"crypto/rand"
	"io"
	"log"
	"path/filepath"
	"strconv"

	"github.com/gorilla/sessions"
	"github.com/mholt/caddy"
	"github.com/mholt/caddy/caddyhttp/httpserver"
	"github.com/upframe/fest/email"
	"github.com/upframe/fest/models"
	"github.com/upframe/fest/pages"
)

var (
	// Store stores the session cookies and help us to handle them
	store *sessions.CookieStore
	// BaseAddress is the base URL to build URLs
	BaseAddress string
	// RootPath is the 'root' directive defined in Caddyfile
	RootPath string
	// TemplatesPath is where the templates are stored
	TemplatesPath string
)

func init() {
	// Regists the caddy middleware
	caddy.RegisterPlugin("fest", caddy.Plugin{
		ServerType: "http",
		Action:     setup,
	})
}

func setup(c *caddy.Controller) error {
	// Gets the base address
	cfg := httpserver.GetConfig(c)

	pages.BaseAddress = cfg.Addr.String()
	pages.Templates = filepath.Clean(cfg.Root+"/templates/") + string(filepath.Separator)
	email.Templates = pages.Templates + "email" + string(filepath.Separator)

	var (
		err                                    error
		smtpUser, smtpPass, smtpHost, smtpPort string
		dbUser, dbPass, dbHost, dbName         string
		dbPort                                 = "3306"
		development                            = false
		keyPairs                               [][]byte
	)

	// Gets the options from the Caddyfile
	for c.Next() {
		for c.NextBlock() {
			switch c.Val() {
			case "db_user":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbUser = c.Val()
			case "db_pass":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbPass = c.Val()
			case "db_host":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbHost = c.Val()
			case "db_port":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbPort = c.Val()
			case "db_name":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbName = c.Val()
			case "smtp_user":
				if !c.NextArg() {
					return c.ArgErr()
				}

				smtpUser = c.Val()
			case "smtp_pass":
				if !c.NextArg() {
					return c.ArgErr()
				}

				smtpPass = c.Val()
			case "smtp_host":
				if !c.NextArg() {
					return c.ArgErr()
				}

				smtpHost = c.Val()
			case "smtp_port":
				if !c.NextArg() {
					return c.ArgErr()
				}

				smtpPort = c.Val()
			case "base_invites":
				if !c.NextArg() {
					return c.ArgErr()
				}

				pages.BaseInvites, err = strconv.Atoi(c.Val())
				if err != nil {
					return err
				}
			case "development":
				development = true
			}
		}
	}

	// Sets up the cookies
	if !development {
		// Generates 5 random key pairs to secure the cookies
		// NOTE: generating this at startup will automatically log out the
		// users when the server is rebooted
		keyPairs = [][]byte{}
		for i := 0; i < 5; i++ {
			keyPairs = append(keyPairs, make([]byte, 32))
			_, err = io.ReadFull(rand.Reader, keyPairs[i])
			if err != nil {
				log.Fatal(err)
			}
		}
	} else {
		keyPairs = [][]byte{[]byte("HEY")}
	}

	// Creates the new cookie session;
	store = sessions.NewCookieStore(keyPairs...)
	store.Options = &sessions.Options{
		Path:     "/",
		MaxAge:   3600 * 3,
		Secure:   cfg.Addr.Scheme == "https",
		HttpOnly: true,
	}
	store.Options.Domain = cfg.Host()

	// Configures the email
	email.InitSMTP(smtpUser, smtpPass, smtpHost, smtpPort)

	// Connects to the database and checks for an error
	err = models.InitDB(dbUser, dbPass, dbHost, dbPort, dbName)
	if err != nil {
		return err
	}

	// Configures PayPal TODO: get from caddyfile
	if pages.InitPayPal(
		"AeMshfG7Rpy6zsw7GCREUBXi897HzkwuVshrf2BmYk2uU7Q4If3ax4AOJtaEqfU8lS-QymRuQYX0R9LL",
		"EIYSLuOY6B19Yqc_2tONbjOdhHxR15w166Pa3JxiEylZZn7ZSWC7SZtpSLbOJcLS9n7fs391hZ9XY6N6",
		development,
	) != nil {
		return err
	}

	// Adds the middleware to Caddy
	mid := func(next httpserver.Handler) httpserver.Handler {
		return Upframe{
			Next: next,
			Root: cfg.Root,
		}
	}

	httpserver.GetConfig(c).AddMiddleware(mid)
	return nil
}
<commit_msg>paypal data in caddyfile<commit_after>package fest

import (
	"crypto/rand"
	"io"
	"log"
	"path/filepath"
	"strconv"

	"github.com/gorilla/sessions"
	"github.com/mholt/caddy"
	"github.com/mholt/caddy/caddyhttp/httpserver"
	"github.com/upframe/fest/email"
	"github.com/upframe/fest/models"
	"github.com/upframe/fest/pages"
)

var (
	// Store stores the session cookies and help us to handle them
	store *sessions.CookieStore
	// BaseAddress is the base URL to build URLs
	BaseAddress string
	// RootPath is the 'root' directive defined in Caddyfile
	RootPath string
	// TemplatesPath is where the templates are stored
	TemplatesPath string
)

func init() {
	// Regists the caddy middleware
	caddy.RegisterPlugin("fest", caddy.Plugin{
		ServerType: "http",
		Action:     setup,
	})
}

func setup(c *caddy.Controller) error {
	// Gets the base address
	cfg := httpserver.GetConfig(c)

	pages.BaseAddress = cfg.Addr.String()
	pages.Templates = filepath.Clean(cfg.Root+"/templates/") + string(filepath.Separator)
	email.Templates = pages.Templates + "email" + string(filepath.Separator)

	var (
		err                                    error
		smtpUser, smtpPass, smtpHost, smtpPort string
		dbUser, dbPass, dbHost, dbName         string
		paypalClient, paypalSecret             string
		dbPort                                 = "3306"
		development                            = false
		keyPairs                               [][]byte
	)

	// Gets the options from the Caddyfile
	for c.Next() {
		for c.NextBlock() {
			switch c.Val() {
			case "db_user":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbUser = c.Val()
			case "db_pass":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbPass = c.Val()
			case "db_host":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbHost = c.Val()
			case "db_port":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbPort = c.Val()
			case "db_name":
				if !c.NextArg() {
					return c.ArgErr()
				}

				dbName = c.Val()
			case "smtp_user":
				if !c.NextArg() {
					return c.ArgErr()
				}

				smtpUser = c.Val()
			case "smtp_pass":
				if !c.NextArg() {
					return c.ArgErr()
				}

				smtpPass = c.Val()
			case "smtp_host":
				if !c.NextArg() {
					return c.ArgErr()
				}

				smtpHost = c.Val()
			case "smtp_port":
				if !c.NextArg() {
					return c.ArgErr()
				}

				smtpPort = c.Val()
			case "base_invites":
				if !c.NextArg() {
					return c.ArgErr()
				}

				pages.BaseInvites, err = strconv.Atoi(c.Val())
				if err != nil {
					return err
				}
			case "paypal_client":
				if !c.NextArg() {
					return c.ArgErr()
				}

				paypalClient = c.Val()
			case "paypal_secret":
				if !c.NextArg() {
					return c.ArgErr()
				}

				paypalSecret = c.Val()
			case "development":
				development = true
			}
		}
	}

	// Sets up the cookies
	if !development {
		// Generates 5 random key pairs to secure the cookies
		// NOTE: generating this at startup will automatically log out the
		// users when the server is rebooted
		keyPairs = [][]byte{}
		for i := 0; i < 5; i++ {
			keyPairs = append(keyPairs, make([]byte, 32))
			_, err = io.ReadFull(rand.Reader, keyPairs[i])
			if err != nil {
				log.Fatal(err)
			}
		}
	} else {
		keyPairs = [][]byte{[]byte("HEY")}
	}

	// Creates the new cookie session;
	store = sessions.NewCookieStore(keyPairs...)
	store.Options = &sessions.Options{
		Path:     "/",
		MaxAge:   3600 * 3,
		Secure:   cfg.Addr.Scheme == "https",
		HttpOnly: true,
	}
	store.Options.Domain = cfg.Host()

	// Configures the email
	email.InitSMTP(smtpUser, smtpPass, smtpHost, smtpPort)

	// Connects to the database and checks for an error
	err = models.InitDB(dbUser, dbPass, dbHost, dbPort, dbName)
	if err != nil {
		return err
	}

	// Configures PayPal
	if pages.InitPayPal(paypalClient, paypalSecret, development) != nil {
		return err
	}

	// Adds the middleware to Caddy
	mid := func(next httpserver.Handler) httpserver.Handler {
		return Upframe{
			Next: next,
			Root: cfg.Root,
		}
	}

	httpserver.GetConfig(c).AddMiddleware(mid)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "net/http"
  "os"
  "io/ioutil"
)

func main(){
  var url string
  fmt.Println(">> url: (Make sure to use http://)")
  fmt.Scanf("%s", &url)
  if url == "" {
    fmt.Println("Please type in an url")
    os.Exit(1)
  }
  fmt.Println("Fetching url...")
  /* fmt.Println("Fetching url...")
  if err = !nil {
    fmt.Println("Something went wrong...")
    os.Exit(1)
  } */
  data, err := http.Get(url)
  if err != nil{
    fmt.Println("Please use http:// in your url")
    os.Exit(1)
  }else{
    defer data.Body.Close()
    contents, err := ioutil.ReadAll(data.Body)
    if err != nil{
      fmt.Println("%s\n", err)
    }
    fmt.Println("%s\n", string(contents))
  }
}
<commit_msg>Moved fetching url message if the url is correct<commit_after>package main

import (
  "fmt"
  "net/http"
  "os"
  "io/ioutil"
)

func main(){
  var url string
  fmt.Println(">> url: (Make sure to use http://)")
  fmt.Scanf("%s", &url)
  if url == "" {
    fmt.Println("Please type in an url")
    os.Exit(1)
  }
  // fmt.Println("Fetching url...")
  /* fmt.Println("Fetching url...")
  if err = !nil {
    fmt.Println("Something went wrong...")
    os.Exit(1)
  } */
  data, err := http.Get(url)
  if err != nil{
    fmt.Println("Please use http:// in your url")
    os.Exit(1)
  }else{
    fmt.Println("Fetching url...")
    defer data.Body.Close()
    contents, err := ioutil.ReadAll(data.Body)
    if err != nil{
      fmt.Println("%s\n", err)
    }
    fmt.Println("%s\n", string(contents))
  }
}
<|endoftext|>


<commit_before>package mcset

import (
	"errors"
	"net"
	"sort"
	"sync"
	"time"

	"github.com/golang/groupcache/consistenthash"
	"github.com/reusee/mmh3"
)

var (
	// ErrNoServers is returned when no servers are configured or available.
	ErrNoServers = errors.New("mcset: no servers configured or available")
)

// A Watcher represents how a serverset.Watch is used so it can be stubbed out for tests.
type Watcher interface {
	Endpoints() []string
	Event() <-chan struct{}
	IsClosed() bool
}

// A MCSet is a wrapper around the serverset.Watch to handle the memcache use case.
// Basically provides some helper functions to pick the servers consistently.
type MCSet struct {
	Watcher

	LastEvent  time.Time
	EventCount int

	// This channel will get an event when zookeeper updates things
	// calling SetEndpoints will not trigger this type of event.
	event chan struct{}

	consistent *consistenthash.Map

	lock      sync.Mutex
	endpoints []string
	addresses map[string]net.Addr
}

// New creates a new memcache server set.
// Can be used to just consistently hash keys to a known set of servers by
// having watch = nil and then calling SetEndpoints with the known set of memcache hosts.
func New(watch Watcher) *MCSet {
	mcset := &MCSet{
		Watcher: watch,
		event:   make(chan struct{}, 1),
	}

	if watch != nil {
		// first time don't trigger an event
		mcset.setEndpoints(watch.Endpoints())

		go func() {
			for {
				select {
				case <-watch.Event():
					mcset.SetEndpoints(watch.Endpoints())
				}

				if watch.IsClosed() {
					break
				}
			}

			watcherClosed()
		}()
	}

	return mcset
}

// for use during testing. Saw this in the net/http standard lib.
var watcherClosed = func() {}

// SetEndpoints sets current list of endpoints. This will override the list
// returned by the serverset. An event by the serverset will override these values.
func (s *MCSet) SetEndpoints(endpoints []string) {
	s.setEndpoints(endpoints)
	s.triggerEvent()
}

func (s *MCSet) setEndpoints(endpoints []string) {
	s.lock.Lock()
	defer s.lock.Unlock()

	addresses := make(map[string]net.Addr)
	for _, e := range endpoints {

		a, err := net.ResolveTCPAddr("tcp", e)
		if err != nil {
			// TODO: if the hostname doesn't resolve, what should we do?
			// panic(err)
		}

		addresses[e] = a
	}

	s.addresses = addresses
	s.endpoints = endpoints

	sort.StringSlice(endpoints).Sort()

	s.consistent = consistenthash.New(150, mmh3.Sum32)
	s.consistent.Add(endpoints...)
}

// Endpoints returns the current endpoints for this service.
// This can be those set via the serverset.Watch or manually via SetEndpoints()
func (s *MCSet) Endpoints() []string {
	return s.endpoints
}

// Event returns the event channel. This channel will get an object
// whenever something changes with the list of endpoints.
// Mostly just a passthrough of the underlying watch event.
func (s *MCSet) Event() <-chan struct{} {
	return s.event
}

// PickServer consistently picks a server from the list.
// Kind of a weird signature but is necessary to satisfy the memcache.ServerSelector interface.
func (s *MCSet) PickServer(key string) (net.Addr, error) {
	if s.consistent == nil {
		return nil, ErrNoServers
	}

	server := s.consistent.Get(key)
	if server == "" {
		return nil, ErrNoServers
	}
	return s.addresses[server], nil
}

// Each runs the function over each server currently in the set.
// Kind of a weird signature but is necessary to satisfy the memcache.ServerSelector interface.
func (s *MCSet) Each(f func(net.Addr) error) error {
	addresses := s.addresses

	for _, a := range addresses {
		if err := f(a); nil != err {
			return err
		}
	}

	return nil
}

// triggerEvent, will queue up something in the Event channel if there isn't already something there.
func (s *MCSet) triggerEvent() {
	s.EventCount++
	s.LastEvent = time.Now()

	select {
	case s.event <- struct{}{}:
	default:
	}
}
<commit_msg>Add logging to mcset<commit_after>package mcset

import (
	"errors"
	"log"
	"net"
	"sort"
	"sync"
	"time"

	"github.com/golang/groupcache/consistenthash"
	"github.com/reusee/mmh3"
)

var (
	// ErrNoServers is returned when no servers are configured or available.
	ErrNoServers = errors.New("mcset: no servers configured or available")

	// DefaultLogger is used by default to print change event messages.
	DefaultLogger Logger = defaultLogger{}
)

// Logger is an interface that can be implemented to provide custom log output.
type Logger interface {
	Printf(string, ...interface{})
}

// A Watcher represents how a serverset.Watch is used so it can be stubbed out for tests.
type Watcher interface {
	Endpoints() []string
	Event() <-chan struct{}
	IsClosed() bool
}

// A MCSet is a wrapper around the serverset.Watch to handle the memcache use case.
// Basically provides some helper functions to pick the servers consistently.
type MCSet struct {
	Watcher

	LastEvent  time.Time
	EventCount int

	Logger Logger

	// This channel will get an event when zookeeper updates things
	// calling SetEndpoints will not trigger this type of event.
	event chan struct{}

	consistent *consistenthash.Map

	lock      sync.Mutex
	endpoints []string
	addresses map[string]net.Addr
}

// New creates a new memcache server set.
// Can be used to just consistently hash keys to a known set of servers by
// having watch = nil and then calling SetEndpoints with the known set of memcache hosts.
func New(watch Watcher) *MCSet {
	mcset := &MCSet{
		Watcher: watch,
		Logger:  DefaultLogger,
		event:   make(chan struct{}, 1),
	}

	if watch != nil {
		// first time don't trigger an event
		mcset.setEndpoints(watch.Endpoints())

		go func() {
			for {
				select {
				case <-watch.Event():
					mcset.SetEndpoints(watch.Endpoints())
				}

				if watch.IsClosed() {
					break
				}
			}

			watcherClosed()
		}()
	}

	return mcset
}

// for use during testing. Saw this in the net/http standard lib.
var watcherClosed = func() {}

// SetEndpoints sets current list of endpoints. This will override the list
// returned by the serverset. An event by the serverset will override these values.
func (s *MCSet) SetEndpoints(endpoints []string) {
	s.setEndpoints(endpoints)
	s.triggerEvent()
}

func (s *MCSet) setEndpoints(endpoints []string) {
	s.lock.Lock()
	defer s.lock.Unlock()

	addresses := make(map[string]net.Addr)
	for _, e := range endpoints {

		a, err := net.ResolveTCPAddr("tcp", e)
		if err != nil {
			// TODO: if the hostname doesn't resolve, what should we do?
			// panic(err)
		}

		addresses[e] = a
	}

	s.addresses = addresses
	s.endpoints = endpoints

	sort.StringSlice(endpoints).Sort()

	s.Logger.Printf("new endpoints for mcset: %v", endpoints)

	s.consistent = consistenthash.New(150, mmh3.Sum32)
	s.consistent.Add(endpoints...)
}

// Endpoints returns the current endpoints for this service.
// This can be those set via the serverset.Watch or manually via SetEndpoints()
func (s *MCSet) Endpoints() []string {
	return s.endpoints
}

// Event returns the event channel. This channel will get an object
// whenever something changes with the list of endpoints.
// Mostly just a passthrough of the underlying watch event.
func (s *MCSet) Event() <-chan struct{} {
	return s.event
}

// PickServer consistently picks a server from the list.
// Kind of a weird signature but is necessary to satisfy the memcache.ServerSelector interface.
func (s *MCSet) PickServer(key string) (net.Addr, error) {
	if s.consistent == nil {
		return nil, ErrNoServers
	}

	server := s.consistent.Get(key)
	if server == "" {
		return nil, ErrNoServers
	}
	return s.addresses[server], nil
}

// Each runs the function over each server currently in the set.
// Kind of a weird signature but is necessary to satisfy the memcache.ServerSelector interface.
func (s *MCSet) Each(f func(net.Addr) error) error {
	addresses := s.addresses

	for _, a := range addresses {
		if err := f(a); nil != err {
			return err
		}
	}

	return nil
}

// triggerEvent, will queue up something in the Event channel if there isn't already something there.
func (s *MCSet) triggerEvent() {
	s.EventCount++
	s.LastEvent = time.Now()

	select {
	case s.event <- struct{}{}:
	default:
	}
}

type defaultLogger struct{}

func (defaultLogger) Printf(format string, a ...interface{}) {
	log.Printf(format, a...)
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"net"
	"time"

	"github.com/nange/easyss/cipherstream"
	"github.com/nange/easyss/utils"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

type state int

const (
	ESTABLISHED state = iota
	FIN_WAIT1
	FIN_WAIT2
	LAST_ACK
	CLOSING
	CLOSE_WAIT
	TIME_WAIT
	CLOSED
)

var stateMap = map[state]string{
	ESTABLISHED: "state: ESTABLISHED",
	FIN_WAIT1:   "state: FIN_WAIT1",
	FIN_WAIT2:   "state: FIN_WAIT2",
	LAST_ACK:    "state: LAST_ACK",
	CLOSING:     "state: CLOSING",
	CLOSE_WAIT:  "state: CLOSE_WAIT",
	TIME_WAIT:   "state: TIME_WAIT",
	CLOSED:      "state: CLOSED",
}

func (s state) String() string {
	if _, ok := stateMap[s]; ok {
		return stateMap[s]
	}
	return "unknown state"
}

type ConnStateFn func(conn io.ReadWriteCloser) *ConnState

type ConnState struct {
	fn    ConnStateFn
	state state
	err   error
	buf   []byte
}

func NewConnState(s state) *ConnState {
	cs := &ConnState{
		state: s,
		buf:   make([]byte, 64),
	}
	statefnMap := map[state]ConnStateFn{
		FIN_WAIT1:  cs.FINWait1,
		FIN_WAIT2:  cs.FINWait2,
		LAST_ACK:   cs.LastACK,
		CLOSING:    cs.Closing,
		CLOSE_WAIT: cs.CloseWait,
		TIME_WAIT:  cs.TimeWait,
		CLOSED:     cs.Closed,
	}
	if statefn, ok := statefnMap[s]; ok {
		cs.fn = statefn
	}
	return cs
}

func (cs *ConnState) FINWait1(conn io.ReadWriteCloser) *ConnState {
	log.Info("start FINWait1 state")
	defer log.Info("end FINWait1 state")

	cs.state = FIN_WAIT1
	fin := utils.NewFINRstStreamHeader()
	_, err := conn.Write(fin)
	if err != nil {
		log.Errorf("conn.Write FIN err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	for {
		_, err = conn.Read(cs.buf)
		log.Infof("FINWAIT1 conn.Read, err:%v", err)
		if err != nil {
			break
		}
	}
	if cipherstream.FINRSTStreamErr(err) {
		cs.fn = cs.Closing
		return cs
	}

	if cipherstream.ACKRSTStreamErr(err) {
		cs.fn = cs.FINWait2
		return cs
	}

	log.Errorf("except get ErrFINRSTStream or ErrACKRSTStream, but get: %v", err)
	cs.err = err
	cs.fn = nil
	return cs
}

func (cs *ConnState) FINWait2(conn io.ReadWriteCloser) *ConnState {
	log.Info("start FINWait2 state")
	defer log.Info("end FINWait2 state")
	cs.state = FIN_WAIT2
	var err error
	for {
		_, err = conn.Read(cs.buf)
		if err != nil {
			break
		}
	}
	if !cipherstream.FINRSTStreamErr(err) {
		log.Errorf("except get ErrFINRSTStream, but get: %+v", err)
		cs.err = err
		cs.fn = nil
		return cs
	}

	ack := utils.NewACKRstStreamHeader()
	_, err = conn.Write(ack)
	if err != nil {
		log.Errorf("conn.Write ACK err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	cs.fn = cs.TimeWait
	return cs
}

func (cs *ConnState) LastACK(conn io.ReadWriteCloser) *ConnState {
	log.Info("start LastACK state")
	defer log.Info("end LastACK state")

	cs.state = LAST_ACK
	fin := utils.NewFINRstStreamHeader()
	_, err := conn.Write(fin)
	if err != nil {
		log.Errorf("conn.Write FIN err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	cs.fn = cs.Closed
	return cs
}

func (cs *ConnState) Closing(conn io.ReadWriteCloser) *ConnState {
	log.Info("start Closing state")
	defer log.Info("end Closing state")

	cs.state = CLOSING
	ack := utils.NewACKRstStreamHeader()
	_, err = conn.Write(ack)
	if err != nil {
		log.Errorf("conn.Write ACK err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	cs.fn = cs.Closed
	return cs
}

func (cs *ConnState) CloseWait(conn io.ReadWriteCloser) *ConnState {
	log.Info("start CloseWait state")
	defer log.Info("end CloseWait state")

	cs.state = CLOSE_WAIT
	ack := utils.NewACKRstStreamHeader()
	_, err := conn.Write(ack)
	if err != nil {
		log.Errorf("conn.Write ack err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	cs.fn = cs.LastACK
	return cs
}

func (cs *ConnState) TimeWait(conn io.ReadWriteCloser) *ConnState {
	log.Info("start TimeWait state")
	defer log.Info("end TimeWait state")

	cs.state = TIME_WAIT
	log.Info("in our TimeWait state, we should end our state machine immediately")

	// if conn is tcp connection, set the deadline to default
	var err error
	if cs, ok := conn.(*cipherstream.CipherStream); ok {
		if c, ok := cs.ReadWriteCloser.(net.Conn); ok {
			err = c.SetDeadline(time.Time{})
			log.Info("set tcp connection deadline to default")
		}
	}
	if err != nil {
		log.Errorf("conn.SetDeadline to default err:%+v", errors.WithStack(err))
		cs.err = err
	}
	cs.fn = nil
	return cs
}

func (cs *ConnState) Closed(conn io.ReadWriteCloser) *ConnState {
	log.Info("start Closed state")
	defer log.Info("end Closed state")

	cs.state = CLOSED
	var err error
	for {
		_, err = conn.Read(cs.buf)
		if err != nil {
			break
		}
	}
	if !cipherstream.ACKRSTStreamErr(err) {
		log.Errorf("except get ErrACKRSTStream, but get: %+v", err)
		cs.err = err
		cs.fn = nil
		return cs
	}

	// if conn is tcp connection, set the deadline to default
	if cs, ok := conn.(*cipherstream.CipherStream); ok {
		if c, ok := cs.ReadWriteCloser.(net.Conn); ok {
			err = c.SetDeadline(time.Time{})
			log.Info("set tcp connection deadline to default")
		}
	}
	if err != nil {
		log.Errorf("conn.SetDeadline to default err:%+v", errors.WithStack(err))
		cs.err = err
	}
	cs.fn = nil

	return cs
}
<commit_msg>fix build<commit_after>package main

import (
	"io"
	"net"
	"time"

	"github.com/nange/easyss/cipherstream"
	"github.com/nange/easyss/utils"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

type state int

const (
	ESTABLISHED state = iota
	FIN_WAIT1
	FIN_WAIT2
	LAST_ACK
	CLOSING
	CLOSE_WAIT
	TIME_WAIT
	CLOSED
)

var stateMap = map[state]string{
	ESTABLISHED: "state: ESTABLISHED",
	FIN_WAIT1:   "state: FIN_WAIT1",
	FIN_WAIT2:   "state: FIN_WAIT2",
	LAST_ACK:    "state: LAST_ACK",
	CLOSING:     "state: CLOSING",
	CLOSE_WAIT:  "state: CLOSE_WAIT",
	TIME_WAIT:   "state: TIME_WAIT",
	CLOSED:      "state: CLOSED",
}

func (s state) String() string {
	if _, ok := stateMap[s]; ok {
		return stateMap[s]
	}
	return "unknown state"
}

type ConnStateFn func(conn io.ReadWriteCloser) *ConnState

type ConnState struct {
	fn    ConnStateFn
	state state
	err   error
	buf   []byte
}

func NewConnState(s state) *ConnState {
	cs := &ConnState{
		state: s,
		buf:   make([]byte, 64),
	}
	statefnMap := map[state]ConnStateFn{
		FIN_WAIT1:  cs.FINWait1,
		FIN_WAIT2:  cs.FINWait2,
		LAST_ACK:   cs.LastACK,
		CLOSING:    cs.Closing,
		CLOSE_WAIT: cs.CloseWait,
		TIME_WAIT:  cs.TimeWait,
		CLOSED:     cs.Closed,
	}
	if statefn, ok := statefnMap[s]; ok {
		cs.fn = statefn
	}
	return cs
}

func (cs *ConnState) FINWait1(conn io.ReadWriteCloser) *ConnState {
	log.Info("start FINWait1 state")
	defer log.Info("end FINWait1 state")

	cs.state = FIN_WAIT1
	fin := utils.NewFINRstStreamHeader()
	_, err := conn.Write(fin)
	if err != nil {
		log.Errorf("conn.Write FIN err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	for {
		_, err = conn.Read(cs.buf)
		log.Infof("FINWAIT1 conn.Read, err:%v", err)
		if err != nil {
			break
		}
	}
	if cipherstream.FINRSTStreamErr(err) {
		cs.fn = cs.Closing
		return cs
	}

	if cipherstream.ACKRSTStreamErr(err) {
		cs.fn = cs.FINWait2
		return cs
	}

	log.Errorf("except get ErrFINRSTStream or ErrACKRSTStream, but get: %v", err)
	cs.err = err
	cs.fn = nil
	return cs
}

func (cs *ConnState) FINWait2(conn io.ReadWriteCloser) *ConnState {
	log.Info("start FINWait2 state")
	defer log.Info("end FINWait2 state")
	cs.state = FIN_WAIT2
	var err error
	for {
		_, err = conn.Read(cs.buf)
		if err != nil {
			break
		}
	}
	if !cipherstream.FINRSTStreamErr(err) {
		log.Errorf("except get ErrFINRSTStream, but get: %+v", err)
		cs.err = err
		cs.fn = nil
		return cs
	}

	ack := utils.NewACKRstStreamHeader()
	_, err = conn.Write(ack)
	if err != nil {
		log.Errorf("conn.Write ACK err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	cs.fn = cs.TimeWait
	return cs
}

func (cs *ConnState) LastACK(conn io.ReadWriteCloser) *ConnState {
	log.Info("start LastACK state")
	defer log.Info("end LastACK state")

	cs.state = LAST_ACK
	fin := utils.NewFINRstStreamHeader()
	_, err := conn.Write(fin)
	if err != nil {
		log.Errorf("conn.Write FIN err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	cs.fn = cs.Closed
	return cs
}

func (cs *ConnState) Closing(conn io.ReadWriteCloser) *ConnState {
	log.Info("start Closing state")
	defer log.Info("end Closing state")

	cs.state = CLOSING
	ack := utils.NewACKRstStreamHeader()
	_, err := conn.Write(ack)
	if err != nil {
		log.Errorf("conn.Write ACK err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	cs.fn = cs.Closed
	return cs
}

func (cs *ConnState) CloseWait(conn io.ReadWriteCloser) *ConnState {
	log.Info("start CloseWait state")
	defer log.Info("end CloseWait state")

	cs.state = CLOSE_WAIT
	ack := utils.NewACKRstStreamHeader()
	_, err := conn.Write(ack)
	if err != nil {
		log.Errorf("conn.Write ack err:%+v", errors.WithStack(err))
		cs.err = err
		cs.fn = nil
		return cs
	}

	cs.fn = cs.LastACK
	return cs
}

func (cs *ConnState) TimeWait(conn io.ReadWriteCloser) *ConnState {
	log.Info("start TimeWait state")
	defer log.Info("end TimeWait state")

	cs.state = TIME_WAIT
	log.Info("in our TimeWait state, we should end our state machine immediately")

	// if conn is tcp connection, set the deadline to default
	var err error
	if cs, ok := conn.(*cipherstream.CipherStream); ok {
		if c, ok := cs.ReadWriteCloser.(net.Conn); ok {
			err = c.SetDeadline(time.Time{})
			log.Info("set tcp connection deadline to default")
		}
	}
	if err != nil {
		log.Errorf("conn.SetDeadline to default err:%+v", errors.WithStack(err))
		cs.err = err
	}
	cs.fn = nil
	return cs
}

func (cs *ConnState) Closed(conn io.ReadWriteCloser) *ConnState {
	log.Info("start Closed state")
	defer log.Info("end Closed state")

	cs.state = CLOSED
	var err error
	for {
		_, err = conn.Read(cs.buf)
		if err != nil {
			break
		}
	}
	if !cipherstream.ACKRSTStreamErr(err) {
		log.Errorf("except get ErrACKRSTStream, but get: %+v", err)
		cs.err = err
		cs.fn = nil
		return cs
	}

	// if conn is tcp connection, set the deadline to default
	if cs, ok := conn.(*cipherstream.CipherStream); ok {
		if c, ok := cs.ReadWriteCloser.(net.Conn); ok {
			err = c.SetDeadline(time.Time{})
			log.Info("set tcp connection deadline to default")
		}
	}
	if err != nil {
		log.Errorf("conn.SetDeadline to default err:%+v", errors.WithStack(err))
		cs.err = err
	}
	cs.fn = nil

	return cs
}
<|endoftext|>