<commit_before>package logtalez

import (
	"fmt"
	"io"
	"log"
	"testing"

	"github.com/zeromq/goczmq"
)

func TestMakeTopicList(t *testing.T) {
	hosts := "host1,host2"
	programs := "program1,program2"
	topics := MakeTopicList(hosts, programs)

	expected := []string{
		"host1.program1",
		"host1.program2",
		"host2.program1",
		"host2.program2",
	}

	for i, expect := range expected {
		if topics[i] != expect {
			t.Errorf("expected topic '%s', got '%s'", expect[i], topics[i])
		}
	}
}

func TestMakeEndpointList(t *testing.T) {
	conns := "tcp://incproc1,tcp://inproc2"
	endpoints := MakeEndpointList(conns)

	expected := []string{
		"tcp://incproc1",
		"tcp://inproc2",
	}

	for i, expect := range expected {
		if endpoints[i] != expect {
			t.Errorf("expected endpoint '%s', got '%s'", expect[i], endpoints[i])
		}
	}
}

func TestNew(t *testing.T) {
	endpoints := []string{"inproc://test1"}
	topics := []string{"topic1", "topic2"}
	servercert := "./example_certs/example_curve_server_cert"
	clientcert := "./example_certs/example_curve_client_cert"

	auth := goczmq.NewAuth()
	defer auth.Destroy()

	clientCert, err := goczmq.NewCertFromFile(clientcert)
	if err != nil {
		t.Fatal(err)
	}
	defer clientCert.Destroy()

	server := goczmq.NewSock(goczmq.Pub)

	defer server.Destroy()
	server.SetZapDomain("global")

	serverCert, err := goczmq.NewCertFromFile(servercert)
	defer serverCert.Destroy()
	if err != nil {
		t.Fatal(err)
	}

	serverCert.Apply(server)
	server.SetCurveServer(1)

	err = auth.Curve("./example_certs/")
	if err != nil {
		t.Fatal(err)
	}

	server.Bind(endpoints[0])

	lt, err := New(endpoints, topics, servercert, clientcert)
	if err != nil {
		t.Error("NewLogTalez failed: %s", err)
	}

	server.SendFrame([]byte("topic1:hello world"), 0)

	buf := make([]byte, 65536)

	n, err := lt.Read(buf)
	if err != io.EOF {
		t.Errorf("expected %s, got %s", io.EOF, err)
	}

	if string(buf[:n]) != "topic1:hello world" {
		t.Errorf("expected 'topic1:hello world', got '%s'", buf[:n])
	}

	server.SendFrame([]byte("topic2:hello again"), 0)

	n, err = lt.Read(buf)
	if err != io.EOF {
		t.Errorf("expected %s, got %s", io.EOF, err)
	}

	if string(buf[:n]) != "topic2:hello again" {
		t.Errorf("expected 'topic2:hello again', got '%s'", buf[:n])
	}
}

func ExampleLogTalez() {
	// endpoints is a list of zeromq rsyslog endpoints to attach to.
	endpoints := []string{"tcp://host1.example.com:24444,tcp://host2.example.com:24444"}

	// topics is a list of topics to subscribe to.
	topics := []string{"host1.ssh,host1.nginx,host2.ssh,host2.kernel"}

	// path to the server public certificate
	serverCertPath := "/home/example_user/.curve/server_cert"

	// path to the client public certificate
	clientCertPath := "/home/example_user/.curve/my_cert"

	// create a new logtalez instance
	lt, err := New(endpoints, topics, serverCertPath, clientCertPath)
	if err != nil {
		log.Fatal(err)
	}
	defer lt.Destroy()

	// logtalez exposes an io.Reader interface.  so, here we create a
	// buffer and read log lines.  currently logtalez returns one log
	// line per call to Read.
	buf := make([]byte, 65536)
	for {
		n, err := lt.Read(buf)
		if err != nil && err != io.EOF {
			panic(err)
		}

		fmt.Println(string(buf[:n]))
	}
}
<commit_msg>fix logtalez test to reflect goczmq io.Reader API changes<commit_after>package logtalez

import (
	"fmt"
	"io"
	"log"
	"testing"

	"github.com/zeromq/goczmq"
)

func TestMakeTopicList(t *testing.T) {
	hosts := "host1,host2"
	programs := "program1,program2"
	topics := MakeTopicList(hosts, programs)

	expected := []string{
		"host1.program1",
		"host1.program2",
		"host2.program1",
		"host2.program2",
	}

	for i, expect := range expected {
		if topics[i] != expect {
			t.Errorf("expected topic '%s', got '%s'", expect[i], topics[i])
		}
	}
}

func TestMakeEndpointList(t *testing.T) {
	conns := "tcp://incproc1,tcp://inproc2"
	endpoints := MakeEndpointList(conns)

	expected := []string{
		"tcp://incproc1",
		"tcp://inproc2",
	}

	for i, expect := range expected {
		if endpoints[i] != expect {
			t.Errorf("expected endpoint '%s', got '%s'", expect[i], endpoints[i])
		}
	}
}

func TestNew(t *testing.T) {
	endpoints := []string{"inproc://test1"}
	topics := []string{"topic1", "topic2"}
	servercert := "./example_certs/example_curve_server_cert"
	clientcert := "./example_certs/example_curve_client_cert"

	auth := goczmq.NewAuth()
	defer auth.Destroy()

	clientCert, err := goczmq.NewCertFromFile(clientcert)
	if err != nil {
		t.Fatal(err)
	}
	defer clientCert.Destroy()

	server := goczmq.NewSock(goczmq.Pub)

	defer server.Destroy()
	server.SetZapDomain("global")

	serverCert, err := goczmq.NewCertFromFile(servercert)
	defer serverCert.Destroy()
	if err != nil {
		t.Fatal(err)
	}

	serverCert.Apply(server)
	server.SetCurveServer(1)

	err = auth.Curve("./example_certs/")
	if err != nil {
		t.Fatal(err)
	}

	server.Bind(endpoints[0])

	lt, err := New(endpoints, topics, servercert, clientcert)
	if err != nil {
		t.Error("NewLogTalez failed: %s", err)
	}

	server.SendFrame([]byte("topic1:hello world"), 0)

	buf := make([]byte, 65536)

	n, err := lt.Read(buf)
	if err != nil {
		t.Errorf("expected nil, got %s", err)
	}

	if string(buf[:n]) != "topic1:hello world" {
		t.Errorf("expected 'topic1:hello world', got '%s'", buf[:n])
	}

	server.SendFrame([]byte("topic2:hello again"), 0)

	n, err = lt.Read(buf)
	if err != nil {
		t.Errorf("expected nil, got %s", err)
	}

	if string(buf[:n]) != "topic2:hello again" {
		t.Errorf("expected 'topic2:hello again', got '%s'", buf[:n])
	}
}

func ExampleLogTalez() {
	// endpoints is a list of zeromq rsyslog endpoints to attach to.
	endpoints := []string{"tcp://host1.example.com:24444,tcp://host2.example.com:24444"}

	// topics is a list of topics to subscribe to.
	topics := []string{"host1.ssh,host1.nginx,host2.ssh,host2.kernel"}

	// path to the server public certificate
	serverCertPath := "/home/example_user/.curve/server_cert"

	// path to the client public certificate
	clientCertPath := "/home/example_user/.curve/my_cert"

	// create a new logtalez instance
	lt, err := New(endpoints, topics, serverCertPath, clientCertPath)
	if err != nil {
		log.Fatal(err)
	}
	defer lt.Destroy()

	// logtalez exposes an io.Reader interface.  so, here we create a
	// buffer and read log lines.  currently logtalez returns one log
	// line per call to Read.
	buf := make([]byte, 65536)
	for {
		n, err := lt.Read(buf)
		if err != nil && err != io.EOF {
			panic(err)
		}

		fmt.Println(string(buf[:n]))
	}
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"log"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/ghodss/yaml"
)

// if help is more than tooLong lines long, send a private message
const tooLong = 14

// Size of QR code
const qrsize = 400

// If this list doesn't match what's registered below,
// you're gonna have a bad time.
var builtIns = []string{
	"builtInhelp",
	"builtInadmin",
	"builtIndump",
	"builtInlogging",
}

func init() {
	RegisterPlugin("builtIndump", PluginHandler{DefaultConfig: dumpConfig, Handler: dump})
	RegisterPlugin("builtInhelp", PluginHandler{DefaultConfig: helpConfig, Handler: help})
	RegisterPlugin("builtInadmin", PluginHandler{DefaultConfig: adminConfig, Handler: admin})
	RegisterPlugin("builtInlogging", PluginHandler{DefaultConfig: logConfig, Handler: logging})
}

/* builtin plugins, like help */

func help(bot *Robot, command string, args ...string) (retval PlugRetVal) {
	if command == "init" {
		return // ignore init
	}
	if command == "info" {
		robot.RLock()
		admins := strings.Join(robot.adminUsers, ", ")
		alias := robot.alias
		robot.RUnlock()
		msg := make([]string, 0, 7)
		msg = append(msg, "Here's some information about my running environment:")
		msg = append(msg, fmt.Sprintf("The hostname for the server I'm running on is: %s", hostName))
		if bot.CheckAdmin() {
			msg = append(msg, fmt.Sprintf("My install directory is: %s", robot.installPath))
			msg = append(msg, fmt.Sprintf("My local configuration directory is: %s", robot.localPath))
		}
		msg = append(msg, fmt.Sprintf("My software version is: Gopherbot %s, commit: %s", Version, commit))
		if alias != 0 {
			msg = append(msg, fmt.Sprintf("My alias is: %s", string(alias)))
		}
		msg = append(msg, fmt.Sprintf("The administrators for this robot are: %s", admins))
		bot.Say(strings.Join(msg, "\n"))
	}
	if command == "help" {
		robot.RLock()
		defer robot.RUnlock()

		var term, helpOutput string
		botSub := `(bot)`
		hasTerm := false
		lineSeparator := "\n\n"

		if len(args) == 1 && len(args[0]) > 0 {
			hasTerm = true
			term = args[0]
			if term == "help" {
				Log(Trace, "Help requested for help, returning")
				return
			}
			Log(Trace, "Help requested for term", term)
		}

		helpLines := make([]string, 0, tooLong)
		currentPlugins.RLock()
		plugins := currentPlugins.p
		currentPlugins.RUnlock()
		for _, plugin := range plugins {
			Log(Trace, fmt.Sprintf("Checking help for plugin %s (term: %s)", plugin.name, term))
			if !hasTerm { // if you ask for help without a term, you just get help for whatever commands are available to you
				if pluginAvailable(bot.User, bot.Channel, plugin) {
					for _, phelp := range plugin.Help {
						for _, helptext := range phelp.Helptext {
							if len(phelp.Keywords) > 0 && phelp.Keywords[0] == "*" {
								// * signifies help that should be prepended
								newSize := tooLong
								if len(helpLines) > newSize {
									newSize += len(helpLines)
								}
								prepend := make([]string, 1, newSize)
								prepend[0] = strings.Replace(helptext, botSub, robot.name, -1)
								helpLines = append(prepend, helpLines...)
							} else {
								helpLines = append(helpLines, strings.Replace(helptext, botSub, robot.name, -1))
							}
						}
					}
				}
			} else { // when there's a search term, give all help for that term, but add (channels: xxx) at the end
				for _, phelp := range plugin.Help {
					for _, keyword := range phelp.Keywords {
						if term == keyword {
							chantext := ""
							if plugin.DirectOnly {
								// Look: the right paren gets added below
								chantext = " (direct message only"
							} else {
								for _, pchan := range plugin.Channels {
									if len(chantext) == 0 {
										chantext += " (channels: " + pchan
									} else {
										chantext += ", " + pchan
									}
								}
							}
							if len(chantext) != 0 {
								chantext += ")"
							}
							for _, helptext := range phelp.Helptext {
								helpLines = append(helpLines, strings.Replace(helptext, botSub, robot.name, -1)+chantext)
							}
						}
					}
				}
			}
		}
		if hasTerm {
			helpOutput = "Command(s) matching keyword: " + term + "\n" + strings.Join(helpLines, lineSeparator)
		}
		switch {
		case len(helpLines) == 0:
			bot.Say("Sorry, bub - I got nothin' for ya'")
		case len(helpLines) > tooLong:
			if len(bot.Channel) > 0 {
				bot.Reply("(the help output was pretty long, so I sent you a private message)")
				if !hasTerm {
					helpOutput = "Command(s) available in channel: " + bot.Channel + "\n" + strings.Join(helpLines, lineSeparator)
				}
			} else {
				if !hasTerm {
					helpOutput = "Command(s) available:" + "\n" + strings.Join(helpLines, lineSeparator)
				}
			}
			bot.SendUserMessage(bot.User, helpOutput)
		default:
			if !hasTerm {
				helpOutput = "Command(s) available:" + "\n" + strings.Join(helpLines, lineSeparator)
			}
			bot.Say(helpOutput)
		}
	}
	return
}

func dump(bot *Robot, command string, args ...string) (retval PlugRetVal) {
	if command == "init" {
		return // ignore init
	}
	currentPlugins.RLock()
	plugins := currentPlugins.p
	currentPlugins.RUnlock()
	robot.RLock()
	defer robot.RUnlock()
	switch command {
	case "robot":
		c, _ := yaml.Marshal(config)
		bot.Fixed().Say(fmt.Sprintf("Here's how I've been configured, irrespective of interactive changes:\n%s", c))
	case "plugdefault":
		if plug, ok := pluginHandlers[args[0]]; ok {
			bot.Fixed().Say(fmt.Sprintf("Here's the default configuration for \"%s\":\n%s", args[0], plug.DefaultConfig))
		} else { // look for an external plugin
			found := false
			for _, plugin := range plugins {
				if args[0] == plugin.name && plugin.pluginType == plugExternal {
					found = true
					if cfg, err := getExtDefCfg(plugin); err == nil {
						bot.Fixed().Say(fmt.Sprintf("Here's the default configuration for \"%s\":\n%s", args[0], *cfg))
					} else {
						bot.Say("I had a problem looking that up - somebody should check my logs")
					}
				}
			}
			if !found {
				bot.Say("Didn't find a plugin named " + args[0])
			}
		}
	case "plugin":
		found := false
		for _, plugin := range plugins {
			if args[0] == plugin.name {
				found = true
				c, _ := yaml.Marshal(plugin)
				bot.Fixed().Say(fmt.Sprintf("%s", c))
			}
		}
		if !found {
			bot.Say("Didn't find a plugin named " + args[0])
		}
	case "list":
		plist := make([]string, 0, len(plugins))
		for _, plugin := range plugins {
			plist = append(plist, plugin.name)
		}
		bot.Say(fmt.Sprintf("Here are the plugins I have configured:\n%s", strings.Join(plist, ", ")))
	}
	return
}

var byebye = []string{
	"Sayonara!",
	"Adios",
	"Hasta la vista!",
	"Later gator!",
}

func logging(bot *Robot, command string, args ...string) (retval PlugRetVal) {
	switch command {
	case "init":
		return
	case "level":
		setLogLevel(logStrToLevel(args[0]))
		bot.Say(fmt.Sprintf("I've adjusted the log level to %s", args[0]))
		Log(Info, fmt.Sprintf("User %s changed logging level to %s", bot.User, args[0]))
	case "show":
		page := 0
		if len(args) == 1 {
			page, _ = strconv.Atoi(args[0])
		}
		lines, wrap := logPage(page)
		if wrap {
			bot.Say("(warning: value too large for pages, wrapped past beginning of log)")
		}
		bot.Fixed().Say(strings.Join(lines, ""))
	case "showlevel":
		l := getLogLevel()
		bot.Say(fmt.Sprintf("My current logging level is: %s", logLevelToStr(l)))
	case "setlines":
		l, _ := strconv.Atoi(args[0])
		set := setLogPageLines(l)
		bot.Say(fmt.Sprintf("Lines per page of log output set to: %d", set))
	}
	return
}

func admin(bot *Robot, command string, args ...string) (retval PlugRetVal) {
	if command == "init" {
		return // ignore init
	}
	if !bot.CheckAdmin() {
		bot.Reply("Sorry, only an admin user can request that")
		return
	}
	switch command {
	case "reload":
		err := loadConfig()
		if err != nil {
			bot.Reply("Error encountered during reload, check the logs")
			Log(Error, fmt.Errorf("Reloading configuration, requested by %s: %v", bot.User, err))
			return
		}
		bot.Reply("Configuration reloaded successfully")
		Log(Info, "Configuration successfully reloaded by a request from:", bot.User)
	case "abort":
		buf := make([]byte, 32768)
		runtime.Stack(buf, true)
		log.Printf("%s", buf)
		time.Sleep(2 * time.Second)
		panic("Abort command issued")
	case "quit":
		pluginsRunning.Done()
		pluginsRunning.Lock()
		pluginsRunning.count--
		pluginsRunning.shuttingDown = true
		if pluginsRunning.count > 0 {
			runningCount := pluginsRunning.count
			pluginsRunning.Unlock()
			bot.Say(fmt.Sprintf("There are still %d plugins running; I'll exit when they all complete, or you can issue an \"abort\" command", runningCount))
		} else {
			pluginsRunning.Unlock()
		}
		// Wait for all plugins to stop running
		pluginsRunning.Wait()
		bot.Reply(bot.RandomString(byebye))
		// Stop the brain after it finishes any current task
		brainQuit()
		Log(Info, "Exiting on administrator command")
		// How long does it _actually_ take for the message to go out?
		time.Sleep(time.Second)
		close(finish)
	}
	return
}
<commit_msg>No help for unavailable plugins<commit_after>package bot

import (
	"fmt"
	"log"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/ghodss/yaml"
)

// if help is more than tooLong lines long, send a private message
const tooLong = 14

// Size of QR code
const qrsize = 400

// If this list doesn't match what's registered below,
// you're gonna have a bad time.
var builtIns = []string{
	"builtInhelp",
	"builtInadmin",
	"builtIndump",
	"builtInlogging",
}

func init() {
	RegisterPlugin("builtIndump", PluginHandler{DefaultConfig: dumpConfig, Handler: dump})
	RegisterPlugin("builtInhelp", PluginHandler{DefaultConfig: helpConfig, Handler: help})
	RegisterPlugin("builtInadmin", PluginHandler{DefaultConfig: adminConfig, Handler: admin})
	RegisterPlugin("builtInlogging", PluginHandler{DefaultConfig: logConfig, Handler: logging})
}

/* builtin plugins, like help */

func help(bot *Robot, command string, args ...string) (retval PlugRetVal) {
	if command == "init" {
		return // ignore init
	}
	if command == "info" {
		robot.RLock()
		admins := strings.Join(robot.adminUsers, ", ")
		alias := robot.alias
		robot.RUnlock()
		msg := make([]string, 0, 7)
		msg = append(msg, "Here's some information about my running environment:")
		msg = append(msg, fmt.Sprintf("The hostname for the server I'm running on is: %s", hostName))
		if bot.CheckAdmin() {
			msg = append(msg, fmt.Sprintf("My install directory is: %s", robot.installPath))
			msg = append(msg, fmt.Sprintf("My local configuration directory is: %s", robot.localPath))
		}
		msg = append(msg, fmt.Sprintf("My software version is: Gopherbot %s, commit: %s", Version, commit))
		if alias != 0 {
			msg = append(msg, fmt.Sprintf("My alias is: %s", string(alias)))
		}
		msg = append(msg, fmt.Sprintf("The administrators for this robot are: %s", admins))
		bot.Say(strings.Join(msg, "\n"))
	}
	if command == "help" {
		robot.RLock()
		defer robot.RUnlock()

		var term, helpOutput string
		botSub := `(bot)`
		hasTerm := false
		lineSeparator := "\n\n"

		if len(args) == 1 && len(args[0]) > 0 {
			hasTerm = true
			term = args[0]
			if term == "help" {
				Log(Trace, "Help requested for help, returning")
				return
			}
			Log(Trace, "Help requested for term", term)
		}

		helpLines := make([]string, 0, tooLong)
		currentPlugins.RLock()
		plugins := currentPlugins.p
		currentPlugins.RUnlock()
		for _, plugin := range plugins {
			if !pluginAvailable(bot.User, bot.Channel, plugin) {
				continue
			}
			Log(Trace, fmt.Sprintf("Checking help for plugin %s (term: %s)", plugin.name, term))
			if !hasTerm { // if you ask for help without a term, you just get help for whatever commands are available to you
				for _, phelp := range plugin.Help {
					for _, helptext := range phelp.Helptext {
						if len(phelp.Keywords) > 0 && phelp.Keywords[0] == "*" {
							// * signifies help that should be prepended
							newSize := tooLong
							if len(helpLines) > newSize {
								newSize += len(helpLines)
							}
							prepend := make([]string, 1, newSize)
							prepend[0] = strings.Replace(helptext, botSub, robot.name, -1)
							helpLines = append(prepend, helpLines...)
						} else {
							helpLines = append(helpLines, strings.Replace(helptext, botSub, robot.name, -1))
						}
					}
				}
			} else { // when there's a search term, give all help for that term, but add (channels: xxx) at the end
				for _, phelp := range plugin.Help {
					for _, keyword := range phelp.Keywords {
						if term == keyword {
							chantext := ""
							if plugin.DirectOnly {
								// Look: the right paren gets added below
								chantext = " (direct message only"
							} else {
								for _, pchan := range plugin.Channels {
									if len(chantext) == 0 {
										chantext += " (channels: " + pchan
									} else {
										chantext += ", " + pchan
									}
								}
							}
							if len(chantext) != 0 {
								chantext += ")"
							}
							for _, helptext := range phelp.Helptext {
								helpLines = append(helpLines, strings.Replace(helptext, botSub, robot.name, -1)+chantext)
							}
						}
					}
				}
			}
		}
		if hasTerm {
			helpOutput = "Command(s) matching keyword: " + term + "\n" + strings.Join(helpLines, lineSeparator)
		}
		switch {
		case len(helpLines) == 0:
			bot.Say("Sorry, bub - I got nothin' for ya'")
		case len(helpLines) > tooLong:
			if len(bot.Channel) > 0 {
				bot.Reply("(the help output was pretty long, so I sent you a private message)")
				if !hasTerm {
					helpOutput = "Command(s) available in channel: " + bot.Channel + "\n" + strings.Join(helpLines, lineSeparator)
				}
			} else {
				if !hasTerm {
					helpOutput = "Command(s) available:" + "\n" + strings.Join(helpLines, lineSeparator)
				}
			}
			bot.SendUserMessage(bot.User, helpOutput)
		default:
			if !hasTerm {
				helpOutput = "Command(s) available:" + "\n" + strings.Join(helpLines, lineSeparator)
			}
			bot.Say(helpOutput)
		}
	}
	return
}

func dump(bot *Robot, command string, args ...string) (retval PlugRetVal) {
	if command == "init" {
		return // ignore init
	}
	currentPlugins.RLock()
	plugins := currentPlugins.p
	currentPlugins.RUnlock()
	robot.RLock()
	defer robot.RUnlock()
	switch command {
	case "robot":
		c, _ := yaml.Marshal(config)
		bot.Fixed().Say(fmt.Sprintf("Here's how I've been configured, irrespective of interactive changes:\n%s", c))
	case "plugdefault":
		if plug, ok := pluginHandlers[args[0]]; ok {
			bot.Fixed().Say(fmt.Sprintf("Here's the default configuration for \"%s\":\n%s", args[0], plug.DefaultConfig))
		} else { // look for an external plugin
			found := false
			for _, plugin := range plugins {
				if args[0] == plugin.name && plugin.pluginType == plugExternal {
					found = true
					if cfg, err := getExtDefCfg(plugin); err == nil {
						bot.Fixed().Say(fmt.Sprintf("Here's the default configuration for \"%s\":\n%s", args[0], *cfg))
					} else {
						bot.Say("I had a problem looking that up - somebody should check my logs")
					}
				}
			}
			if !found {
				bot.Say("Didn't find a plugin named " + args[0])
			}
		}
	case "plugin":
		found := false
		for _, plugin := range plugins {
			if args[0] == plugin.name {
				found = true
				c, _ := yaml.Marshal(plugin)
				bot.Fixed().Say(fmt.Sprintf("%s", c))
			}
		}
		if !found {
			bot.Say("Didn't find a plugin named " + args[0])
		}
	case "list":
		plist := make([]string, 0, len(plugins))
		for _, plugin := range plugins {
			plist = append(plist, plugin.name)
		}
		bot.Say(fmt.Sprintf("Here are the plugins I have configured:\n%s", strings.Join(plist, ", ")))
	}
	return
}

var byebye = []string{
	"Sayonara!",
	"Adios",
	"Hasta la vista!",
	"Later gator!",
}

func logging(bot *Robot, command string, args ...string) (retval PlugRetVal) {
	switch command {
	case "init":
		return
	case "level":
		setLogLevel(logStrToLevel(args[0]))
		bot.Say(fmt.Sprintf("I've adjusted the log level to %s", args[0]))
		Log(Info, fmt.Sprintf("User %s changed logging level to %s", bot.User, args[0]))
	case "show":
		page := 0
		if len(args) == 1 {
			page, _ = strconv.Atoi(args[0])
		}
		lines, wrap := logPage(page)
		if wrap {
			bot.Say("(warning: value too large for pages, wrapped past beginning of log)")
		}
		bot.Fixed().Say(strings.Join(lines, ""))
	case "showlevel":
		l := getLogLevel()
		bot.Say(fmt.Sprintf("My current logging level is: %s", logLevelToStr(l)))
	case "setlines":
		l, _ := strconv.Atoi(args[0])
		set := setLogPageLines(l)
		bot.Say(fmt.Sprintf("Lines per page of log output set to: %d", set))
	}
	return
}

func admin(bot *Robot, command string, args ...string) (retval PlugRetVal) {
	if command == "init" {
		return // ignore init
	}
	if !bot.CheckAdmin() {
		bot.Reply("Sorry, only an admin user can request that")
		return
	}
	switch command {
	case "reload":
		err := loadConfig()
		if err != nil {
			bot.Reply("Error encountered during reload, check the logs")
			Log(Error, fmt.Errorf("Reloading configuration, requested by %s: %v", bot.User, err))
			return
		}
		bot.Reply("Configuration reloaded successfully")
		Log(Info, "Configuration successfully reloaded by a request from:", bot.User)
	case "abort":
		buf := make([]byte, 32768)
		runtime.Stack(buf, true)
		log.Printf("%s", buf)
		time.Sleep(2 * time.Second)
		panic("Abort command issued")
	case "quit":
		pluginsRunning.Done()
		pluginsRunning.Lock()
		pluginsRunning.count--
		pluginsRunning.shuttingDown = true
		if pluginsRunning.count > 0 {
			runningCount := pluginsRunning.count
			pluginsRunning.Unlock()
			bot.Say(fmt.Sprintf("There are still %d plugins running; I'll exit when they all complete, or you can issue an \"abort\" command", runningCount))
		} else {
			pluginsRunning.Unlock()
		}
		// Wait for all plugins to stop running
		pluginsRunning.Wait()
		bot.Reply(bot.RandomString(byebye))
		// Stop the brain after it finishes any current task
		brainQuit()
		Log(Info, "Exiting on administrator command")
		// How long does it _actually_ take for the message to go out?
		time.Sleep(time.Second)
		close(finish)
	}
	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package uid

import (
	"encoding/json"
	"io"

	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/log"
	"github.com/mutecomm/mute/util"
	"github.com/mutecomm/mute/util/times"
)

// A SessionAnchor contains the keys for perfect forward secrecy.
type SessionAnchor struct {
	MIXADDRESS string     // fully qualified address of mix to use as last hop to user
	NYMADDRESS string     // a valid NymAddress
	PFKEYS     []KeyEntry // for ephemeral/forward secure key agreement
}

type contents struct {
	VERSION           string // the protocol version
	MSGCOUNT          uint64 // must increase for each message of the same type and user
	NOTAFTER          uint64 // time after which the key(s) should not be used anymore
	NOTBEFORE         uint64 // time before which the key(s) should not be used yet
	FALLBACK          bool   // determines if the key may serve as a fallback key
	SIGKEYHASH        string // SHA512(UIDMessage.UIDContent.SIGKEY.HASH)
	REPOURI           string // URI of the corresponding KeyInit repository
	SESSIONANCHOR     string // SESSIONANCHOR = AES256_CTR(key=UIDMessage.UIDContent.SIGKEY.HASH, SessionAnchor)
	SESSIONANCHORHASH string // before encryption
}

// A KeyInit message contains short-term keys.
type KeyInit struct {
	CONTENTS  contents
	SIGNATURE string // signature of contents by UIDMessage.UIDContent.SIGKEY
}

// The number of seconds the NOTAFTER field of a KeyInit message can be in the
// future.
const MaxNotAfter = uint64(90 * 24 * 60 * 60) // 90 days

// NewJSONKeyInit returns a new KeyInit message initialized with the parameters
// given in the JSON byte array.
func NewJSONKeyInit(keyInit []byte) (*KeyInit, error) {
	var ki KeyInit
	if err := json.Unmarshal(keyInit, &ki); err != nil {
		return nil, err
	}
	return &ki, nil
}

func (sa *SessionAnchor) json() []byte {
	return marshalSorted(sa)
}

// PrivateKey returns the base64 encoded private signature key of
// session anchor.
func (sa *SessionAnchor) PrivateKey() string {
	return base64.Encode(sa.PFKEYS[0].PrivateKey32()[:])
}

// SetPrivateKey sets the private key to the given base64 encoded privkey
// string.
func (sa *SessionAnchor) SetPrivateKey(privkey string) error {
	key, err := base64.Decode(privkey)
	if err != nil {
		return err
	}
	return sa.PFKEYS[0].setPrivateKey(key)
}

// KeyEntry returns the KeyEntry of the SessionAnchor for the given function.
func (sa *SessionAnchor) KeyEntry(function string) (*KeyEntry, error) {
	for _, ke := range sa.PFKEYS {
		if ke.FUNCTION == function {
			return &ke, nil
		}
	}
	log.Error(ErrKeyEntryNotFound)
	return nil, ErrKeyEntryNotFound
}

// NymAddress returns the nymaddress of the SessionAnchor.
func (sa *SessionAnchor) NymAddress() string {
	return sa.NYMADDRESS
}

func (c *contents) json() []byte {
	return marshalSorted(c)
}

// JSON encodes KeyInit as a JSON string according to the specification.
func (ki *KeyInit) JSON() []byte {
	return marshalSorted(ki)
}

// MsgCount returns the message count of the KeyInit message.
func (ki *KeyInit) MsgCount() uint64 {
	return ki.CONTENTS.MSGCOUNT
}

// SigKeyHash returns the signature key hash of the KeyInit message.
func (ki *KeyInit) SigKeyHash() string {
	return ki.CONTENTS.SIGKEYHASH
}

// SessionAnchor returns the decrypted and verified session anchor for KeyInit.
func (ki *KeyInit) SessionAnchor(sigPubKey string) (*SessionAnchor, error) {
	// SIGKEYHASH corresponds to the SIGKEY of the Identity
	pubKey, err := base64.Decode(sigPubKey)
	if err != nil {
		return nil, err
	}
	keyHash := cipher.SHA512(pubKey)
	if ki.CONTENTS.SIGKEYHASH != base64.Encode(cipher.SHA512(keyHash)) {
		log.Error(ErrWrongSigKeyHash)
		return nil, ErrWrongSigKeyHash
	}
	// verify that SESSIONANCHORHASH matches decrypted SESSIONANCHOR
	enc, err := base64.Decode(ki.CONTENTS.SESSIONANCHOR)
	if err != nil {
		return nil, err
	}
	txt := cipher.AES256CTRDecrypt(keyHash[:32], enc)
	var sa SessionAnchor
	if err := json.Unmarshal(txt, &sa); err != nil {
		return nil, err
	}
	if ki.CONTENTS.SESSIONANCHORHASH != base64.Encode(cipher.SHA512(sa.json())) {
		log.Error(ErrSessionAnchor)
		return nil, ErrSessionAnchor
	}
	return &sa, nil
}

// KeyEntryECDHE25519 returns the decrypted and verified ECDHE25519 KeyEntry
// for KeyInit.
func (ki *KeyInit) KeyEntryECDHE25519(sigPubKey string) (*KeyEntry, error) {
	sa, err := ki.SessionAnchor(sigPubKey)
	if err != nil {
		return nil, err
	}
	ke, err := sa.KeyEntry("ECDHE25519")
	if err != nil {
		return nil, err
	}
	return ke, err
}

// Verify verifies that the KeyInit is valid.
func (ki *KeyInit) Verify(keyInitRepositoryURIs []string, sigPubKey string) error {
	// The REPOURI points to this KeyInit Repository
	if !util.ContainsString(keyInitRepositoryURIs, ki.CONTENTS.REPOURI) {
		log.Error(ErrRepoURI)
		return ErrRepoURI
	}

	// verify that SESSIONANCHORHASH matches decrypted SESSIONANCHOR
	if _, err := ki.SessionAnchor(sigPubKey); err != nil {
		return err
	}

	// NOTAFTER and NOTBEFORE are valid
	if ki.CONTENTS.NOTBEFORE >= ki.CONTENTS.NOTAFTER {
		log.Error(ErrInvalidTimes)
		return ErrInvalidTimes
	}
	// not expired
	if ki.CONTENTS.NOTAFTER < uint64(times.Now()) {
		log.Error(ErrExpired)
		return ErrExpired
	}

	// SIGNATURE was made with UIDMessage.UIDContent.SIGKEY over Contents
	var ed25519Key cipher.Ed25519Key
	sig, err := base64.Decode(ki.SIGNATURE)
	if err != nil {
		return err
	}
	pubKey, err := base64.Decode(sigPubKey)
	if err != nil {
		return err
	}
	// create ed25519 key
	ed25519Key.SetPublicKey(pubKey)
	// verify self-signature
	if !ed25519Key.Verify(ki.CONTENTS.json(), sig) {
		log.Error(ErrInvalidKeyInitSig)
		return ErrInvalidKeyInitSig
	}

	return nil
}

// Sign signs the KeyInit message and returns the signature.
func (ki *KeyInit) Sign(sigKey *cipher.Ed25519Key) string {
	return base64.Encode(sigKey.Sign(ki.JSON()))
}

// VerifySrvSig verifies the signature with the srvPubKey.
func (ki *KeyInit) VerifySrvSig(signature, srvPubKey string) error {
	var ed25519Key cipher.Ed25519Key
	// get server-signature
	sig, err := base64.Decode(signature)
	if err != nil {
		return err
	}
	// create ed25519 key
	pubKey, err := base64.Decode(srvPubKey)
	if err != nil {
		return err
	}
	ed25519Key.SetPublicKey(pubKey)
	// verify server-signature
	if !ed25519Key.Verify(ki.JSON(), sig) {
		log.Error(ErrInvalidSrvSig)
		return ErrInvalidSrvSig
	}
	return nil
}

func (msg *Message) sessionAnchor(
	key []byte,
	mixaddress, nymaddress string,
	rand io.Reader,
) (sessionAnchor, sessionAnchorHash, pubKeyHash, privateKey string, err error) {
	var sa SessionAnchor
	sa.MIXADDRESS = mixaddress
	sa.NYMADDRESS = nymaddress
	sa.PFKEYS = make([]KeyEntry, 1)
	if err := sa.PFKEYS[0].InitDHKey(rand); err != nil {
		return "", "", "", "", err
	}
	jsn := sa.json()
	hash := cipher.SHA512(jsn)
	// SESSIONANCHOR = AES256_CTR(key=UIDMessage.UIDContent.SIGKEY.HASH, SessionAnchor)
	enc := base64.Encode(cipher.AES256CTREncrypt(key[:32], jsn, rand))
	return enc, base64.Encode(hash), sa.PFKEYS[0].HASH, base64.Encode(sa.PFKEYS[0].PrivateKey32()[:]), nil
}

// KeyInit returns a new KeyInit message for the given UID message. It also
// returns the pubKeyHash and privateKey for convenient further use.
// msgcount must increase for each message of the same type and user.
// notafter is the unixtime after which the key(s) should not be used anymore.
// notbefore is the unixtime before which the key(s) should not be used yet.
// fallback determines if the key may serve as a fallback key.
// repoURI is URI of the corresponding KeyInit repository.
// Necessary randomness is read from rand.
func (msg *Message) KeyInit(
	msgcount, notafter, notbefore uint64,
	fallback bool,
	repoURI, mixaddress, nymaddress string,
	rand io.Reader,
) (ki *KeyInit, pubKeyHash, privateKey string, err error) {
	var keyInit KeyInit
	// time checks
	if notbefore >= notafter {
		log.Error(ErrInvalidTimes)
		return nil, "", "", ErrInvalidTimes
	}
	if notafter < uint64(times.Now()) {
		log.Error(ErrExpired)
		return nil, "", "", ErrExpired
	}
	if notafter > uint64(times.Now())+MaxNotAfter {
		log.Error(ErrFuture)
		return nil, "", "", ErrFuture
	}
	// init
	keyInit.CONTENTS.VERSION = ProtocolVersion
	keyInit.CONTENTS.MSGCOUNT = msgcount
	keyInit.CONTENTS.NOTAFTER = notafter
	keyInit.CONTENTS.NOTBEFORE = notbefore
	keyInit.CONTENTS.FALLBACK = fallback
	keyHash, err := base64.Decode(msg.UIDContent.SIGKEY.HASH)
	if err != nil {
		return nil, "", "", err
	}
	keyInit.CONTENTS.SIGKEYHASH = base64.Encode(cipher.SHA512(keyHash))

	// make sure REPOURIS is set to the first REPOURI of UIDContent.REPOURIS
	// TODO: support different KeyInit repository configurations
	if repoURI != msg.UIDContent.REPOURIS[0] {
		return nil, "", "",
			log.Error("uri: repoURI differs from msg.UIDContent.REPOURIS[0]")
	}
	keyInit.CONTENTS.REPOURI = repoURI

	// create SessionAnchor
	sa, sah, pubKeyHash, privateKey, err := msg.sessionAnchor(keyHash,
		mixaddress, nymaddress, rand)
	if err != nil {
		return nil, "", "", err
	}
	keyInit.CONTENTS.SESSIONANCHOR = sa
	keyInit.CONTENTS.SESSIONANCHORHASH = sah
	// sign KeyInit: the content doesn't have to be hashed, because Ed25519 is
	// already taking care of that.
	sig := msg.UIDContent.SIGKEY.ed25519Key.Sign(keyInit.CONTENTS.json())
	keyInit.SIGNATURE = base64.Encode(sig)
	ki = &keyInit
	return
}

// Check that the content of KeyInit is consistent with it's version.
func (ki *KeyInit) Check() error {
	// we only support version 1.0 at this stage
	if ki.CONTENTS.VERSION != "1.0" {
		return log.Errorf("uid: unknown CONTENTS.VERSION: %s",
			ki.CONTENTS.VERSION)
	}
	return nil
}
<commit_msg>uid: ki.CONTENTS -> ki.Contents (as in spec)<commit_after>// Copyright (c) 2015 Mute Communications Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package uid

import (
	"encoding/json"
	"io"

	"github.com/mutecomm/mute/cipher"
	"github.com/mutecomm/mute/encode/base64"
	"github.com/mutecomm/mute/log"
	"github.com/mutecomm/mute/util"
	"github.com/mutecomm/mute/util/times"
)

// A SessionAnchor contains the keys for perfect forward secrecy.
type SessionAnchor struct {
	MIXADDRESS string     // fully qualified address of mix to use as last hop to user
	NYMADDRESS string     // a valid NymAddress
	PFKEYS     []KeyEntry // for ephemeral/forward secure key agreement
}

type contents struct {
	VERSION           string // the protocol version
	MSGCOUNT          uint64 // must increase for each message of the same type and user
	NOTAFTER          uint64 // time after which the key(s) should not be used anymore
	NOTBEFORE         uint64 // time before which the key(s) should not be used yet
	FALLBACK          bool   // determines if the key may serve as a fallback key
	SIGKEYHASH        string // SHA512(UIDMessage.UIDContent.SIGKEY.HASH)
	REPOURI           string // URI of the corresponding KeyInit repository
	SESSIONANCHOR     string // SESSIONANCHOR = AES256_CTR(key=UIDMessage.UIDContent.SIGKEY.HASH, SessionAnchor)
	SESSIONANCHORHASH string // before encryption
}

// A KeyInit message contains short-term keys.
type KeyInit struct {
	Contents  contents
	SIGNATURE string // signature of contents by UIDMessage.UIDContent.SIGKEY
}

// The number of seconds the NOTAFTER field of a KeyInit message can be in the
// future.
const MaxNotAfter = uint64(90 * 24 * 60 * 60) // 90 days

// NewJSONKeyInit returns a new KeyInit message initialized with the parameters
// given in the JSON byte array.
func NewJSONKeyInit(keyInit []byte) (*KeyInit, error) {
	var ki KeyInit
	if err := json.Unmarshal(keyInit, &ki); err != nil {
		return nil, err
	}
	return &ki, nil
}

func (sa *SessionAnchor) json() []byte {
	return marshalSorted(sa)
}

// PrivateKey returns the base64 encoded private signature key of
// session anchor.
func (sa *SessionAnchor) PrivateKey() string {
	return base64.Encode(sa.PFKEYS[0].PrivateKey32()[:])
}

// SetPrivateKey sets the private key to the given base64 encoded privkey
// string.
func (sa *SessionAnchor) SetPrivateKey(privkey string) error {
	key, err := base64.Decode(privkey)
	if err != nil {
		return err
	}
	return sa.PFKEYS[0].setPrivateKey(key)
}

// KeyEntry returns the KeyEntry of the SessionAnchor for the given function.
func (sa *SessionAnchor) KeyEntry(function string) (*KeyEntry, error) {
	for _, ke := range sa.PFKEYS {
		if ke.FUNCTION == function {
			return &ke, nil
		}
	}
	log.Error(ErrKeyEntryNotFound)
	return nil, ErrKeyEntryNotFound
}

// NymAddress returns the nymaddress of the SessionAnchor.
func (sa *SessionAnchor) NymAddress() string {
	return sa.NYMADDRESS
}

func (c *contents) json() []byte {
	return marshalSorted(c)
}

// JSON encodes KeyInit as a JSON string according to the specification.
func (ki *KeyInit) JSON() []byte {
	return marshalSorted(ki)
}

// MsgCount returns the message count of the KeyInit message.
func (ki *KeyInit) MsgCount() uint64 {
	return ki.Contents.MSGCOUNT
}

// SigKeyHash returns the signature key hash of the KeyInit message.
func (ki *KeyInit) SigKeyHash() string {
	return ki.Contents.SIGKEYHASH
}

// SessionAnchor returns the decrypted and verified session anchor for KeyInit.
func (ki *KeyInit) SessionAnchor(sigPubKey string) (*SessionAnchor, error) {
	// SIGKEYHASH corresponds to the SIGKEY of the Identity
	pubKey, err := base64.Decode(sigPubKey)
	if err != nil {
		return nil, err
	}
	keyHash := cipher.SHA512(pubKey)
	if ki.Contents.SIGKEYHASH != base64.Encode(cipher.SHA512(keyHash)) {
		log.Error(ErrWrongSigKeyHash)
		return nil, ErrWrongSigKeyHash
	}
	// verify that SESSIONANCHORHASH matches decrypted SESSIONANCHOR
	enc, err := base64.Decode(ki.Contents.SESSIONANCHOR)
	if err != nil {
		return nil, err
	}
	txt := cipher.AES256CTRDecrypt(keyHash[:32], enc)
	var sa SessionAnchor
	if err := json.Unmarshal(txt, &sa); err != nil {
		return nil, err
	}
	if ki.Contents.SESSIONANCHORHASH != base64.Encode(cipher.SHA512(sa.json())) {
		log.Error(ErrSessionAnchor)
		return nil, ErrSessionAnchor
	}
	return &sa, nil
}

// KeyEntryECDHE25519 returns the decrypted and verified ECDHE25519 KeyEntry
// for KeyInit.
func (ki *KeyInit) KeyEntryECDHE25519(sigPubKey string) (*KeyEntry, error) {
	sa, err := ki.SessionAnchor(sigPubKey)
	if err != nil {
		return nil, err
	}
	ke, err := sa.KeyEntry("ECDHE25519")
	if err != nil {
		return nil, err
	}
	return ke, err
}

// Verify verifies that the KeyInit is valid.
func (ki *KeyInit) Verify(keyInitRepositoryURIs []string, sigPubKey string) error {
	// The REPOURI points to this KeyInit Repository
	if !util.ContainsString(keyInitRepositoryURIs, ki.Contents.REPOURI) {
		log.Error(ErrRepoURI)
		return ErrRepoURI
	}

	// verify that SESSIONANCHORHASH matches decrypted SESSIONANCHOR
	if _, err := ki.SessionAnchor(sigPubKey); err != nil {
		return err
	}

	// NOTAFTER and NOTBEFORE are valid
	if ki.Contents.NOTBEFORE >= ki.Contents.NOTAFTER {
		log.Error(ErrInvalidTimes)
		return ErrInvalidTimes
	}
	// not expired
	if ki.Contents.NOTAFTER < uint64(times.Now()) {
		log.Error(ErrExpired)
		return ErrExpired
	}

	// SIGNATURE was made with UIDMessage.UIDContent.SIGKEY over Contents
	var ed25519Key cipher.Ed25519Key
	sig, err := base64.Decode(ki.SIGNATURE)
	if err != nil {
		return err
	}
	pubKey, err := base64.Decode(sigPubKey)
	if err != nil {
		return err
	}
	// create ed25519 key
	ed25519Key.SetPublicKey(pubKey)
	// verify self-signature
	if !ed25519Key.Verify(ki.Contents.json(), sig) {
		log.Error(ErrInvalidKeyInitSig)
		return ErrInvalidKeyInitSig
	}

	return nil
}

// Sign signs the KeyInit message and returns the signature.
func (ki *KeyInit) Sign(sigKey *cipher.Ed25519Key) string {
	return base64.Encode(sigKey.Sign(ki.JSON()))
}

// VerifySrvSig verifies the signature with the srvPubKey.
func (ki *KeyInit) VerifySrvSig(signature, srvPubKey string) error {
	var ed25519Key cipher.Ed25519Key
	// get server-signature
	sig, err := base64.Decode(signature)
	if err != nil {
		return err
	}
	// create ed25519 key
	pubKey, err := base64.Decode(srvPubKey)
	if err != nil {
		return err
	}
	ed25519Key.SetPublicKey(pubKey)
	// verify server-signature
	if !ed25519Key.Verify(ki.JSON(), sig) {
		log.Error(ErrInvalidSrvSig)
		return ErrInvalidSrvSig
	}
	return nil
}

func (msg *Message) sessionAnchor(
	key []byte,
	mixaddress, nymaddress string,
	rand io.Reader,
) (sessionAnchor, sessionAnchorHash, pubKeyHash, privateKey string, err error) {
	var sa SessionAnchor
	sa.MIXADDRESS = mixaddress
	sa.NYMADDRESS = nymaddress
	sa.PFKEYS = make([]KeyEntry, 1)
	if err := sa.PFKEYS[0].InitDHKey(rand); err != nil {
		return "", "", "", "", err
	}
	jsn := sa.json()
	hash := cipher.SHA512(jsn)
	// SESSIONANCHOR = AES256_CTR(key=UIDMessage.UIDContent.SIGKEY.HASH, SessionAnchor)
	enc := base64.Encode(cipher.AES256CTREncrypt(key[:32], jsn, rand))
	return enc, base64.Encode(hash), sa.PFKEYS[0].HASH, base64.Encode(sa.PFKEYS[0].PrivateKey32()[:]), nil
}

// KeyInit returns a new KeyInit message for the given UID message. It also
// returns the pubKeyHash and privateKey for convenient further use.
// msgcount must increase for each message of the same type and user.
// notafter is the unixtime after which the key(s) should not be used anymore.
// notbefore is the unixtime before which the key(s) should not be used yet.
// fallback determines if the key may serve as a fallback key.
// repoURI is URI of the corresponding KeyInit repository.
// Necessary randomness is read from rand.
func (msg *Message) KeyInit(
	msgcount, notafter, notbefore uint64,
	fallback bool,
	repoURI, mixaddress, nymaddress string,
	rand io.Reader,
) (ki *KeyInit, pubKeyHash, privateKey string, err error) {
	var keyInit KeyInit
	// time checks
	if notbefore >= notafter {
		log.Error(ErrInvalidTimes)
		return nil, "", "", ErrInvalidTimes
	}
	if notafter < uint64(times.Now()) {
		log.Error(ErrExpired)
		return nil, "", "", ErrExpired
	}
	if notafter > uint64(times.Now())+MaxNotAfter {
		log.Error(ErrFuture)
		return nil, "", "", ErrFuture
	}
	// init
	keyInit.Contents.VERSION = ProtocolVersion
	keyInit.Contents.MSGCOUNT = msgcount
	keyInit.Contents.NOTAFTER = notafter
	keyInit.Contents.NOTBEFORE = notbefore
	keyInit.Contents.FALLBACK = fallback
	keyHash, err := base64.Decode(msg.UIDContent.SIGKEY.HASH)
	if err != nil {
		return nil, "", "", err
	}
	keyInit.Contents.SIGKEYHASH = base64.Encode(cipher.SHA512(keyHash))

	// make sure REPOURIS is set to the first REPOURI of UIDContent.REPOURIS
	// TODO: support different KeyInit repository configurations
	if repoURI != msg.UIDContent.REPOURIS[0] {
		return nil, "", "",
			log.Error("uri: repoURI differs from msg.UIDContent.REPOURIS[0]")
	}
	keyInit.Contents.REPOURI = repoURI

	// create SessionAnchor
	sa, sah, pubKeyHash, privateKey, err := msg.sessionAnchor(keyHash,
		mixaddress, nymaddress, rand)
	if err != nil {
		return nil, "", "", err
	}
	keyInit.Contents.SESSIONANCHOR = sa
	keyInit.Contents.SESSIONANCHORHASH = sah
	// sign KeyInit: the content doesn't have to be hashed, because Ed25519 is
	// already taking care of that.
	sig := msg.UIDContent.SIGKEY.ed25519Key.Sign(keyInit.Contents.json())
	keyInit.SIGNATURE = base64.Encode(sig)
	ki = &keyInit
	return
}

// Check that the content of KeyInit is consistent with it's version.
func (ki *KeyInit) Check() error {
	// we only support version 1.0 at this stage
	if ki.Contents.VERSION != "1.0" {
		return log.Errorf("uid: unknown Contents.VERSION: %s",
			ki.Contents.VERSION)
	}
	return nil
}
<|endoftext|>


<commit_before>package cast

import (
	"math"

	"testing"
)

func TestUint64FromUint8(t *testing.T) {

	tests := []struct{
		Value uint8
	}{
		{
			Value: 0,
		},
		{
			Value: 1,
		},
		{
			Value: math.MaxUint8,
		},
	}

	{
		const numRand = 20
		for i:=0; i<numRand; i++ {
			test := struct{
				Value uint8
			}{
				Value: uint8(randomness.Int63n(math.MaxUint8)),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {

		x, err := Uint64(test.Value)
		if nil != err {
			t.Errorf("For test #%d, did not expect an error, but actually got one: (%T) %v", testNumber, err, err)
			continue
		}

		y := uint8(x)

		if expected, actual := test.Value, y; expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v.", testNumber, expected, actual)
			continue
		}
	}
}

func TestUint64FromUint16(t *testing.T) {

	tests := []struct{
		Value uint16
	}{
		{
			Value: 0,
		},
		{
			Value: 1,
		},
		{
			Value: math.MaxUint16,
		},
	}

	{
		const numRand = 20
		for i:=0; i<numRand; i++ {
			test := struct{
				Value uint16
			}{
				Value: uint16(randomness.Int63n(math.MaxUint16)),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {

		x, err := Uint64(test.Value)
		if nil != err {
			t.Errorf("For test #%d, did not expect an error, but actually got one: (%T) %v", testNumber, err, err)
			continue
		}

		y := uint16(x)

		if expected, actual := test.Value, y; expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v.", testNumber, expected, actual)
			continue
		}
	}
}

func TestUint64FromUint32(t *testing.T) {

	tests := []struct{
		Value uint32
	}{
		{
			Value: 0,
		},
		{
			Value: 1,
		},
		{
			Value: math.MaxUint32,
		},
	}

	{
		const numRand = 20
		for i:=0; i<numRand; i++ {
			test := struct{
				Value uint32
			}{
				Value: uint32(randomness.Int63n(math.MaxUint32)),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {

		x, err := Uint64(test.Value)
		if nil != err {
			t.Errorf("For test #%d, did not expect an error, but actually got one: (%T) %v", testNumber, err, err)
			continue
		}

		y := uint32(x)

		if expected, actual := test.Value, y; expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v.", testNumber, expected, actual)
			continue
		}
	}
}

func TestUint64FromUint64(t *testing.T) {

	tests := []struct{
		Value uint64
	}{
		{
			Value: 0,
		},
		{
			Value: 1,
		},
		{
			Value: math.MaxUint64,
		},
	}

	{
		const numRand = 20
		for i:=0; i<numRand; i++ {
			test := struct{
				Value uint64
			}{
				Value: (uint64(randomness.Uint32()) << 32) + uint64(randomness.Uint32()),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {

		x, err := Uint64(test.Value)
		if nil != err {
			t.Errorf("For test #%d, did not expect an error, but actually got one: (%T) %v", testNumber, err, err)
			continue
		}

		y := uint64(x)

		if expected, actual := test.Value, y; expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v.", testNumber, expected, actual)
			continue
		}
	}
}
<commit_msg>cosmetic change. moved some tests in file.<commit_after>package cast

import (
	"math"

	"testing"
)

func TestUint64FromUint64(t *testing.T) {

	tests := []struct{
		Value uint64
	}{
		{
			Value: 0,
		},
		{
			Value: 1,
		},
		{
			Value: math.MaxUint64,
		},
	}

	{
		const numRand = 20
		for i:=0; i<numRand; i++ {
			test := struct{
				Value uint64
			}{
				Value: (uint64(randomness.Uint32()) << 32) + uint64(randomness.Uint32()),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {

		x, err := Uint64(test.Value)
		if nil != err {
			t.Errorf("For test #%d, did not expect an error, but actually got one: (%T) %v", testNumber, err, err)
			continue
		}

		y := uint64(x)

		if expected, actual := test.Value, y; expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v.", testNumber, expected, actual)
			continue
		}
	}
}

func TestUint64FromUint32(t *testing.T) {

	tests := []struct{
		Value uint32
	}{
		{
			Value: 0,
		},
		{
			Value: 1,
		},
		{
			Value: math.MaxUint32,
		},
	}

	{
		const numRand = 20
		for i:=0; i<numRand; i++ {
			test := struct{
				Value uint32
			}{
				Value: uint32(randomness.Int63n(math.MaxUint32)),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {

		x, err := Uint64(test.Value)
		if nil != err {
			t.Errorf("For test #%d, did not expect an error, but actually got one: (%T) %v", testNumber, err, err)
			continue
		}

		y := uint32(x)

		if expected, actual := test.Value, y; expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v.", testNumber, expected, actual)
			continue
		}
	}
}

func TestUint64FromUint16(t *testing.T) {

	tests := []struct{
		Value uint16
	}{
		{
			Value: 0,
		},
		{
			Value: 1,
		},
		{
			Value: math.MaxUint16,
		},
	}

	{
		const numRand = 20
		for i:=0; i<numRand; i++ {
			test := struct{
				Value uint16
			}{
				Value: uint16(randomness.Int63n(math.MaxUint16)),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {

		x, err := Uint64(test.Value)
		if nil != err {
			t.Errorf("For test #%d, did not expect an error, but actually got one: (%T) %v", testNumber, err, err)
			continue
		}

		y := uint16(x)

		if expected, actual := test.Value, y; expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v.", testNumber, expected, actual)
			continue
		}
	}
}

func TestUint64FromUint8(t *testing.T) {

	tests := []struct{
		Value uint8
	}{
		{
			Value: 0,
		},
		{
			Value: 1,
		},
		{
			Value: math.MaxUint8,
		},
	}

	{
		const numRand = 20
		for i:=0; i<numRand; i++ {
			test := struct{
				Value uint8
			}{
				Value: uint8(randomness.Int63n(math.MaxUint8)),
			}
			tests = append(tests, test)
		}
	}

	for testNumber, test := range tests {

		x, err := Uint64(test.Value)
		if nil != err {
			t.Errorf("For test #%d, did not expect an error, but actually got one: (%T) %v", testNumber, err, err)
			continue
		}

		y := uint8(x)

		if expected, actual := test.Value, y; expected != actual {
			t.Errorf("For test #%d, expected %v, but actually got %v.", testNumber, expected, actual)
			continue
		}
	}
}
<|endoftext|>


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

package main

import (
	"fmt"
	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	_ "image/png"
	"log"
	"math/rand"
)

const (
	screenWidth  = 320
	screenHeight = 240
)

var (
	ebitenImage       *ebiten.Image
	ebitenImageWidth  = 0
	ebitenImageHeight = 0
)

type Sprite struct {
	image *ebiten.Image
	x     int
	y     int
	vx    int
	vy    int
}

func (s *Sprite) Update() {
	s.x += s.vx
	s.y += s.vy
	if s.x < 0 {
		s.x = -s.x
		s.vx = -s.vx
	}
	if s.y < 0 {
		s.y = -s.y
		s.vy = -s.vy
	}
	w, h := s.image.Size()
	if screenWidth <= s.x+w {
		s.x = 2*(screenWidth-w) - s.x
		s.vx = -s.vx
	}
	if screenHeight <= s.y+h {
		s.y = 2*(screenHeight-h) - s.y
		s.vy = -s.vy
	}
}

type Sprites []*Sprite

func (s Sprites) Update() {
	for _, sprite := range s {
		sprite.Update()
	}
}

func (s Sprites) Len() int {
	return len(s)
}

func (s Sprites) Dst(i int) (x0, y0, x1, y1 int) {
	ss := s[i]
	return ss.x, ss.y, ss.x + ebitenImageWidth, ss.y + ebitenImageHeight
}

func (s Sprites) Src(i int) (x0, y0, x1, y1 int) {
	return 0, 0, ebitenImageWidth, ebitenImageHeight
}

var sprites = make(Sprites, 10000)

func update(screen *ebiten.Image) error {
	sprites.Update()
	op := &ebiten.DrawImageOptions{
		ImageParts: sprites,
	}
	op.ColorM.Scale(1.0, 1.0, 1.0, 0.5)
	if err := screen.DrawImage(ebitenImage, op); err != nil {
		return err
	}
	ebitenutil.DebugPrint(screen, fmt.Sprintf("FPS: %0.2f\nNum of sprites: %d", ebiten.CurrentFPS(), sprites.Len()))
	return nil
}

func main() {
	var err error
	ebitenImage, _, err = ebitenutil.NewImageFromFile("images/ebiten.png", ebiten.FilterNearest)
	if err != nil {
		log.Fatal(err)
	}
	ebitenImageWidth, ebitenImageHeight = ebitenImage.Size()
	for i, _ := range sprites {
		w, h := ebitenImage.Size()
		x, y := rand.Intn(screenWidth-w), rand.Intn(screenHeight-h)
		vx, vy := 2*rand.Intn(2)-1, 2*rand.Intn(2)-1
		sprites[i] = &Sprite{
			image: ebitenImage,
			x:     x,
			y:     y,
			vx:    vx,
			vy:    vy,
		}
	}
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Sprites (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>example/sprites: Add UI to adjust the number of sprites (#87)<commit_after>// Copyright 2015 Hajime Hoshi
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	_ "image/png"
	"log"
	"math/rand"
)

const (
	screenWidth  = 320
	screenHeight = 240
)

var (
	ebitenImage       *ebiten.Image
	ebitenImageWidth  = 0
	ebitenImageHeight = 0
)

type Sprite struct {
	image *ebiten.Image
	x     int
	y     int
	vx    int
	vy    int
}

func (s *Sprite) Update() {
	s.x += s.vx
	s.y += s.vy
	if s.x < 0 {
		s.x = -s.x
		s.vx = -s.vx
	}
	if s.y < 0 {
		s.y = -s.y
		s.vy = -s.vy
	}
	w, h := s.image.Size()
	if screenWidth <= s.x+w {
		s.x = 2*(screenWidth-w) - s.x
		s.vx = -s.vx
	}
	if screenHeight <= s.y+h {
		s.y = 2*(screenHeight-h) - s.y
		s.vy = -s.vy
	}
}

type Sprites struct {
	sprites []*Sprite
	num     int
}

func (s Sprites) Update() {
	for _, sprite := range s.sprites {
		sprite.Update()
	}
}

func (s Sprites) Len() int {
	return s.num
}

func (s Sprites) Dst(i int) (x0, y0, x1, y1 int) {
	if s.num <= i {
		return 0, 0, 0, 0
	}
	ss := s.sprites[i]
	return ss.x, ss.y, ss.x + ebitenImageWidth, ss.y + ebitenImageHeight
}

func (s Sprites) Src(i int) (x0, y0, x1, y1 int) {
	if s.num <= i {
		return 0, 0, 0, 0
	}
	return 0, 0, ebitenImageWidth, ebitenImageHeight
}

const (
	MinSprites = 0
	MaxSprites = 10000
)

var sprites = &Sprites{make([]*Sprite, MaxSprites), 500}

func update(screen *ebiten.Image) error {
	if ebiten.IsKeyPressed(ebiten.KeyLeft) {
		sprites.num -= 20
		if sprites.num < MinSprites {
			sprites.num = MinSprites
		}
	}
	if ebiten.IsKeyPressed(ebiten.KeyRight) {
		sprites.num += 20
		if MaxSprites < sprites.num {
			sprites.num = MaxSprites
		}
	}
	sprites.Update()

	op := &ebiten.DrawImageOptions{
		ImageParts: sprites,
	}
	op.ColorM.Scale(1.0, 1.0, 1.0, 0.5)
	if err := screen.DrawImage(ebitenImage, op); err != nil {
		return err
	}
	msg := fmt.Sprintf(`FPS: %0.2f
Num of sprites: %d
Press <- or -> to change the number of sprites`, ebiten.CurrentFPS(), sprites.Len())
	ebitenutil.DebugPrint(screen, msg)
	return nil
}

func main() {
	var err error
	ebitenImage, _, err = ebitenutil.NewImageFromFile("images/ebiten.png", ebiten.FilterNearest)
	if err != nil {
		log.Fatal(err)
	}
	ebitenImageWidth, ebitenImageHeight = ebitenImage.Size()
	for i, _ := range sprites.sprites {
		w, h := ebitenImage.Size()
		x, y := rand.Intn(screenWidth-w), rand.Intn(screenHeight-h)
		vx, vy := 2*rand.Intn(2)-1, 2*rand.Intn(2)-1
		sprites.sprites[i] = &Sprite{
			image: ebitenImage,
			x:     x,
			y:     y,
			vx:    vx,
			vy:    vy,
		}
	}
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Sprites (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


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

package web

import (
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/control-center/serviced/proxy"

	"github.com/control-center/serviced/coordinator/client"
	"github.com/control-center/serviced/dao"
	domainService "github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/zzk"
	"github.com/control-center/serviced/zzk/registry"
	"github.com/control-center/serviced/zzk/service"
	"github.com/zenoss/glog"
)

var (
	allportsLock sync.RWMutex
	allports     map[string]chan int // map of port number to channel that destroys the server
	cpDao        dao.ControlPlane
)

func init() {
	allports = make(map[string]chan int)
}

func disablePort(publicEndpointKey service.PublicEndpointKey) {
	// remove the port from our local cache
	delete(allports, publicEndpointKey.Name())

	// find the endpoint that matches this port number for this service (there will only be 1)
	var myService domainService.Service
	var myEndpoint domainService.ServiceEndpoint
	var unused int
	cpDao.GetService(publicEndpointKey.ServiceID(), &myService)
	for _, endpoint := range myService.Endpoints {
		for _, endpointPort := range endpoint.PortList {
			if endpointPort.PortAddr == publicEndpointKey.Name() {
				myEndpoint = endpoint
			}
		}
	}

	// disable port
	myService.EnablePort(myEndpoint.Name, publicEndpointKey.Name(), false)
	cpDao.UpdateService(myService, &unused)
}

func (sc *ServiceConfig) ServePublicPorts(shutdown <-chan (interface{}), dao dao.ControlPlane) {
	cpDao = dao
	go sc.syncAllPublicPorts(shutdown)
}

func (sc *ServiceConfig) CreatePublicPortServer(publicEndpointKey service.PublicEndpointKey, stopChan <-chan int, shutdown <-chan (interface{})) error {
	port := publicEndpointKey.Name()
	listener, err := net.Listen("tcp", port)
	stopChans := []chan bool{}
	if err != nil {
		glog.Errorf("Could not setup TCP listener - %s", err)
		disablePort(publicEndpointKey)
		return err
	}
	glog.Infof("Listening on port %s", port)

	go func() {
		for {
			// accept connection on public port
			localConn, err := listener.Accept()
			if err != nil {
				glog.V(1).Infof("Stopping accept on port %s", port)
				return
			}

			// lookup remote endpoint for this public port
			pepEPInfo, err := sc.getPublicEndpoint(fmt.Sprintf("%s-%d", publicEndpointKey.Name(), int(publicEndpointKey.Type())))
			if err != nil {
				glog.Errorf("%s", err)
			}

			// setup remote connection
			remotePort := fmt.Sprintf("%s:%d", pepEPInfo.privateIP, pepEPInfo.epPort)
			remoteAddr, err := net.ResolveTCPAddr("tcp", remotePort)
			if err != nil {
				glog.Errorf("Cannot resolve remote address - %s: %s", remotePort, err)
				continue
			} else {
				glog.Infof("Resolved remote address - %s", remotePort)
			}

			remoteConn, err := net.DialTCP("tcp", nil, remoteAddr)
			if err != nil {
				glog.Errorf("%s", err)
				continue
			}

			connStopChan := make(chan bool)
			stopChans = append(stopChans, connStopChan)
			if err != nil {
				for _, c := range stopChans {
					c <- true
				}
				disablePort(publicEndpointKey)
				listener.Close()
			}

			// serve proxied requests/responses
			go proxy.ProxyLoop(localConn, remoteConn, connStopChan)
		}
	}()

	go func() {
		// Wait for shutdown, then kill all your connections
		<-stopChan
		for _, c := range stopChans {
			c <- true
		}

		disablePort(publicEndpointKey)
		listener.Close()
		glog.Infof("Closed port %s", port)
		return
	}()

	return nil
}

func (sc *ServiceConfig) syncAllPublicPorts(shutdown <-chan interface{}) error {
	rootConn, err := zzk.GetLocalConnection("/")
	if err != nil {
		glog.Errorf("syncAllPublicPorts - Error getting root zk connection: %v", err)
		return err
	}

	cancelChan := make(chan interface{})
	syncPorts := func(conn client.Connection, parentPath string, childIDs ...string) {
		glog.V(1).Infof("syncPorts STARTING for parentPath:%s childIDs:%v", parentPath, childIDs)

		// start all servers that have been not started and enabled
		newPorts := make(map[string]chan int)
		for _, pepID := range childIDs {
			publicEndpointKey := service.PublicEndpointKey(pepID)
			if publicEndpointKey.Type() == registry.EPTypePort && publicEndpointKey.IsEnabled() {
				port := publicEndpointKey.Name()
				stopChan, running := allports[port]

				if !running {
					// recently enabled port - port should be opened
					stopChan = make(chan int)
					if err := sc.CreatePublicPortServer(publicEndpointKey, stopChan, shutdown); err != nil {
						continue
					}
				}

				newPorts[port] = stopChan
			}
		}

		// stop all servers that have been deleted or disabled
		for port, stopChan := range allports {
			_, found := newPorts[port]
			if !found {
				stopChan <- 0
				close(stopChan)
			}
		}

		//lock for as short a time as possible
		allportsLock.Lock()
		defer allportsLock.Unlock()
		allports = newPorts
		glog.V(1).Infof("allports: %+v", allports)
	}

	for {
		zkServicePEPService := service.ZKServicePublicEndpoints
		glog.V(1).Infof("Running registry.WatchChildren for zookeeper path: %s", zkServicePEPService)
		err := registry.WatchChildren(rootConn, zkServicePEPService, cancelChan, syncPorts, pepWatchError)
		if err != nil {
			glog.V(1).Infof("Will retry in 10 seconds to WatchChildren(%s) due to error: %v", zkServicePEPService, err)
			<-time.After(time.Second * 10)
			continue
		}
		select {
		case <-shutdown:
			close(cancelChan)
			return nil
		default:
		}
	}
}
<commit_msg>Added port and publicendpoint key to error log<commit_after>// Copyright 2014 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package web

import (
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/control-center/serviced/proxy"

	"github.com/control-center/serviced/coordinator/client"
	"github.com/control-center/serviced/dao"
	domainService "github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/zzk"
	"github.com/control-center/serviced/zzk/registry"
	"github.com/control-center/serviced/zzk/service"
	"github.com/zenoss/glog"
)

var (
	allportsLock sync.RWMutex
	allports     map[string]chan int // map of port number to channel that destroys the server
	cpDao        dao.ControlPlane
)

func init() {
	allports = make(map[string]chan int)
}

func disablePort(publicEndpointKey service.PublicEndpointKey) {
	// remove the port from our local cache
	delete(allports, publicEndpointKey.Name())

	// find the endpoint that matches this port number for this service (there will only be 1)
	var myService domainService.Service
	var myEndpoint domainService.ServiceEndpoint
	var unused int
	cpDao.GetService(publicEndpointKey.ServiceID(), &myService)
	for _, endpoint := range myService.Endpoints {
		for _, endpointPort := range endpoint.PortList {
			if endpointPort.PortAddr == publicEndpointKey.Name() {
				myEndpoint = endpoint
			}
		}
	}

	// disable port
	myService.EnablePort(myEndpoint.Name, publicEndpointKey.Name(), false)
	cpDao.UpdateService(myService, &unused)
}

func (sc *ServiceConfig) ServePublicPorts(shutdown <-chan (interface{}), dao dao.ControlPlane) {
	cpDao = dao
	go sc.syncAllPublicPorts(shutdown)
}

func (sc *ServiceConfig) CreatePublicPortServer(publicEndpointKey service.PublicEndpointKey, stopChan <-chan int, shutdown <-chan (interface{})) error {
	port := publicEndpointKey.Name()
	listener, err := net.Listen("tcp", port)
	stopChans := []chan bool{}
	if err != nil {
		glog.Errorf("Could not setup TCP listener for port %s for public endpoint %s - %s", port, publicEndpointKey, err)
		disablePort(publicEndpointKey)
		return err
	}
	glog.Infof("Listening on port %s", port)

	go func() {
		for {
			// accept connection on public port
			localConn, err := listener.Accept()
			if err != nil {
				glog.V(1).Infof("Stopping accept on port %s", port)
				return
			}

			// lookup remote endpoint for this public port
			pepEPInfo, err := sc.getPublicEndpoint(fmt.Sprintf("%s-%d", publicEndpointKey.Name(), int(publicEndpointKey.Type())))
			if err != nil {
				glog.Errorf("%s", err)
			}

			// setup remote connection
			remotePort := fmt.Sprintf("%s:%d", pepEPInfo.privateIP, pepEPInfo.epPort)
			remoteAddr, err := net.ResolveTCPAddr("tcp", remotePort)
			if err != nil {
				glog.Errorf("Cannot resolve remote address - %s: %s", remotePort, err)
				continue
			} else {
				glog.Infof("Resolved remote address - %s", remotePort)
			}

			remoteConn, err := net.DialTCP("tcp", nil, remoteAddr)
			if err != nil {
				glog.Errorf("%s", err)
				continue
			}

			connStopChan := make(chan bool)
			stopChans = append(stopChans, connStopChan)
			if err != nil {
				for _, c := range stopChans {
					c <- true
				}
				disablePort(publicEndpointKey)
				listener.Close()
			}

			// serve proxied requests/responses
			go proxy.ProxyLoop(localConn, remoteConn, connStopChan)
		}
	}()

	go func() {
		// Wait for shutdown, then kill all your connections
		<-stopChan
		for _, c := range stopChans {
			c <- true
		}

		disablePort(publicEndpointKey)
		listener.Close()
		glog.Infof("Closed port %s", port)
		return
	}()

	return nil
}

func (sc *ServiceConfig) syncAllPublicPorts(shutdown <-chan interface{}) error {
	rootConn, err := zzk.GetLocalConnection("/")
	if err != nil {
		glog.Errorf("syncAllPublicPorts - Error getting root zk connection: %v", err)
		return err
	}

	cancelChan := make(chan interface{})
	syncPorts := func(conn client.Connection, parentPath string, childIDs ...string) {
		glog.V(1).Infof("syncPorts STARTING for parentPath:%s childIDs:%v", parentPath, childIDs)

		// start all servers that have been not started and enabled
		newPorts := make(map[string]chan int)
		for _, pepID := range childIDs {
			publicEndpointKey := service.PublicEndpointKey(pepID)
			if publicEndpointKey.Type() == registry.EPTypePort && publicEndpointKey.IsEnabled() {
				port := publicEndpointKey.Name()
				stopChan, running := allports[port]

				if !running {
					// recently enabled port - port should be opened
					stopChan = make(chan int)
					if err := sc.CreatePublicPortServer(publicEndpointKey, stopChan, shutdown); err != nil {
						continue
					}
				}

				newPorts[port] = stopChan
			}
		}

		// stop all servers that have been deleted or disabled
		for port, stopChan := range allports {
			_, found := newPorts[port]
			if !found {
				stopChan <- 0
				close(stopChan)
			}
		}

		//lock for as short a time as possible
		allportsLock.Lock()
		defer allportsLock.Unlock()
		allports = newPorts
		glog.V(1).Infof("allports: %+v", allports)
	}

	for {
		zkServicePEPService := service.ZKServicePublicEndpoints
		glog.V(1).Infof("Running registry.WatchChildren for zookeeper path: %s", zkServicePEPService)
		err := registry.WatchChildren(rootConn, zkServicePEPService, cancelChan, syncPorts, pepWatchError)
		if err != nil {
			glog.V(1).Infof("Will retry in 10 seconds to WatchChildren(%s) due to error: %v", zkServicePEPService, err)
			<-time.After(time.Second * 10)
			continue
		}
		select {
		case <-shutdown:
			close(cancelChan)
			return nil
		default:
		}
	}
}
<|endoftext|>


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

// Package websocket implements a client and server for the WebSocket protocol.
// The protocol is defined at http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol
package websocket

import (
	"bufio"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"
)

const (
	ProtocolVersionHixie75   = -75
	ProtocolVersionHixie76   = -76
	ProtocolVersionHybi00    = 0
	ProtocolVersionHybi08    = 8
	ProtocolVersionHybi13    = 13
	ProtocolVersionHybi      = ProtocolVersionHybi13
	SupportedProtocolVersion = "13, 8"

	ContinuationFrame = 0
	TextFrame         = 1
	BinaryFrame       = 2
	CloseFrame        = 8
	PingFrame         = 9
	PongFrame         = 10
	UnknownFrame      = 255
)

// WebSocket protocol errors.
type ProtocolError struct {
	ErrorString string
}

func (err *ProtocolError) Error() string { return err.ErrorString }

var (
	ErrBadProtocolVersion   = &ProtocolError{"bad protocol version"}
	ErrBadScheme            = &ProtocolError{"bad scheme"}
	ErrBadStatus            = &ProtocolError{"bad status"}
	ErrBadUpgrade           = &ProtocolError{"missing or bad upgrade"}
	ErrBadWebSocketOrigin   = &ProtocolError{"missing or bad WebSocket-Origin"}
	ErrBadWebSocketLocation = &ProtocolError{"missing or bad WebSocket-Location"}
	ErrBadWebSocketProtocol = &ProtocolError{"missing or bad WebSocket-Protocol"}
	ErrBadWebSocketVersion  = &ProtocolError{"missing or bad WebSocket Version"}
	ErrChallengeResponse    = &ProtocolError{"mismatch challenge/response"}
	ErrBadFrame             = &ProtocolError{"bad frame"}
	ErrBadFrameBoundary     = &ProtocolError{"not on frame boundary"}
	ErrNotWebSocket         = &ProtocolError{"not websocket protocol"}
	ErrBadRequestMethod     = &ProtocolError{"bad method"}
	ErrNotSupported         = &ProtocolError{"not supported"}
)

// Addr is an implementation of net.Addr for WebSocket.
type Addr struct {
	*url.URL
}

// Network returns the network type for a WebSocket, "websocket".
func (addr *Addr) Network() string { return "websocket" }

// Config is a WebSocket configuration
type Config struct {
	// A WebSocket server address.
	Location *url.URL

	// A Websocket client origin.
	Origin *url.URL

	// WebSocket subprotocols.
	Protocol []string

	// WebSocket protocol version.
	Version int

	// TLS config for secure WebSocket (wss).
	TlsConfig *tls.Config

	handshakeData map[string]string
}

// serverHandshaker is an interface to handle WebSocket server side handshake.
type serverHandshaker interface {
	// ReadHandshake reads handshake request message from client.
	// Returns http response code and error if any.
	ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err error)

	// AcceptHandshake accepts the client handshake request and sends
	// handshake response back to client.
	AcceptHandshake(buf *bufio.Writer) (err error)

	// NewServerConn creates a new WebSocket connection.
	NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
}

// frameReader is an interface to read a WebSocket frame.
type frameReader interface {
	// Reader is to read payload of the frame.
	io.Reader

	// PayloadType returns payload type.
	PayloadType() byte

	// HeaderReader returns a reader to read header of the frame.
	HeaderReader() io.Reader

	// TrailerReader returns a reader to read trailer of the frame.
	// If it returns nil, there is no trailer in the frame.
	TrailerReader() io.Reader

	// Len returns total length of the frame, including header and trailer.
	Len() int
}

// frameReaderFactory is an interface to creates new frame reader.
type frameReaderFactory interface {
	NewFrameReader() (r frameReader, err error)
}

// frameWriter is an interface to write a WebSocket frame.
type frameWriter interface {
	// Writer is to write playload of the frame.
	io.WriteCloser
}

// frameWriterFactory is an interface to create new frame writer.
type frameWriterFactory interface {
	NewFrameWriter(payloadType byte) (w frameWriter, err error)
}

type frameHandler interface {
	HandleFrame(frame frameReader) (r frameReader, err error)
	WriteClose(status int) (err error)
}

// Conn represents a WebSocket connection.
type Conn struct {
	config  *Config
	request *http.Request

	buf *bufio.ReadWriter
	rwc io.ReadWriteCloser

	rio sync.Mutex
	frameReaderFactory
	frameReader

	wio sync.Mutex
	frameWriterFactory

	frameHandler
	PayloadType        byte
	defaultCloseStatus int
}

// Read implements the io.Reader interface:
// it reads data of a frame from the WebSocket connection.
// if msg is not large enough for the frame data, it fills the msg and next Read
// will read the rest of the frame data.
// it reads Text frame or Binary frame.
func (ws *Conn) Read(msg []byte) (n int, err error) {
	ws.rio.Lock()
	defer ws.rio.Unlock()
again:
	if ws.frameReader == nil {
		frame, err := ws.frameReaderFactory.NewFrameReader()
		if err != nil {
			return 0, err
		}
		ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
		if err != nil {
			return 0, err
		}
		if ws.frameReader == nil {
			goto again
		}
	}
	n, err = ws.frameReader.Read(msg)
	if err == io.EOF {
		if trailer := ws.frameReader.TrailerReader(); trailer != nil {
			io.Copy(ioutil.Discard, trailer)
		}
		ws.frameReader = nil
		goto again
	}
	return n, err
}

// Write implements the io.Writer interface:
// it writes data as a frame to the WebSocket connection.
func (ws *Conn) Write(msg []byte) (n int, err error) {
	ws.wio.Lock()
	defer ws.wio.Unlock()
	w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
	if err != nil {
		return 0, err
	}
	n, err = w.Write(msg)
	w.Close()
	if err != nil {
		return n, err
	}
	return n, err
}

// Close implements the io.Closer interface.
func (ws *Conn) Close() error {
	err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
	if err != nil {
		return err
	}
	return ws.rwc.Close()
}

func (ws *Conn) IsClientConn() bool { return ws.request == nil }
func (ws *Conn) IsServerConn() bool { return ws.request != nil }

// LocalAddr returns the WebSocket Origin for the connection for client, or
// the WebSocket location for server.
func (ws *Conn) LocalAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Origin}
	}
	return &Addr{ws.config.Location}
}

// RemoteAddr returns the WebSocket location for the connection for client, or
// the Websocket Origin for server.
func (ws *Conn) RemoteAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Location}
	}
	return &Addr{ws.config.Origin}
}

var errSetDeadline = errors.New("websocket: cannot set deadline: not using a net.Conn")

// SetDeadline sets the connection's network read & write deadlines.
func (ws *Conn) SetDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetDeadline(t)
	}
	return errSetDeadline
}

// SetReadDeadline sets the connection's network read deadline.
func (ws *Conn) SetReadDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetReadDeadline(t)
	}
	return errSetDeadline
}

// SetWriteDeadline sets the connection's network write deadline.
func (ws *Conn) SetWriteDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetWriteDeadline(t)
	}
	return errSetDeadline
}

// Config returns the WebSocket config.
func (ws *Conn) Config() *Config { return ws.config }

// Request returns the http request upgraded to the WebSocket.
// It is nil for client side.
func (ws *Conn) Request() *http.Request { return ws.request }

// Codec represents a symmetric pair of functions that implement a codec.
type Codec struct {
	Marshal   func(v interface{}) (data []byte, payloadType byte, err error)
	Unmarshal func(data []byte, payloadType byte, v interface{}) (err error)
}

// Send sends v marshaled by cd.Marshal as single frame to ws.
func (cd Codec) Send(ws *Conn, v interface{}) (err error) {
	if err != nil {
		return err
	}
	data, payloadType, err := cd.Marshal(v)
	if err != nil {
		return err
	}
	ws.wio.Lock()
	defer ws.wio.Unlock()
	w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
	_, err = w.Write(data)
	w.Close()
	return err
}

// Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores in v.
func (cd Codec) Receive(ws *Conn, v interface{}) (err error) {
	ws.rio.Lock()
	defer ws.rio.Unlock()
	if ws.frameReader != nil {
		_, err = io.Copy(ioutil.Discard, ws.frameReader)
		if err != nil {
			return err
		}
		ws.frameReader = nil
	}
again:
	frame, err := ws.frameReaderFactory.NewFrameReader()
	if err != nil {
		return err
	}
	frame, err = ws.frameHandler.HandleFrame(frame)
	if err != nil {
		return err
	}
	if frame == nil {
		goto again
	}
	payloadType := frame.PayloadType()
	data, err := ioutil.ReadAll(frame)
	if err != nil {
		return err
	}
	return cd.Unmarshal(data, payloadType, v)
}

func marshal(v interface{}) (msg []byte, payloadType byte, err error) {
	switch data := v.(type) {
	case string:
		return []byte(data), TextFrame, nil
	case []byte:
		return data, BinaryFrame, nil
	}
	return nil, UnknownFrame, ErrNotSupported
}

func unmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	switch data := v.(type) {
	case *string:
		*data = string(msg)
		return nil
	case *[]byte:
		*data = msg
		return nil
	}
	return ErrNotSupported
}

/*
Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
To send/receive text frame, use string type.
To send/receive binary frame, use []byte type.

Trivial usage:

	import "websocket"

	// receive text frame
	var message string
	websocket.Message.Receive(ws, &message)

	// send text frame
	message = "hello"
	websocket.Message.Send(ws, message)

	// receive binary frame
	var data []byte
	websocket.Message.Receive(ws, &data)

	// send binary frame
	data = []byte{0, 1, 2}
	websocket.Message.Send(ws, data)

*/
var Message = Codec{marshal, unmarshal}

func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err error) {
	msg, err = json.Marshal(v)
	return msg, TextFrame, err
}

func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	return json.Unmarshal(msg, v)
}

/*
JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.

Trival usage:

	import "websocket"

	type T struct {
		Msg string
		Count int
	}

	// receive JSON type T
	var data T
	websocket.JSON.Receive(ws, &data)

	// send JSON type T
	websocket.JSON.Send(ws, data)
*/
var JSON = Codec{jsonMarshal, jsonUnmarshal}
<commit_msg>go.net/websocket: fix error check in Send<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package websocket implements a client and server for the WebSocket protocol.
// The protocol is defined at http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol
package websocket

import (
	"bufio"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"
)

const (
	ProtocolVersionHixie75   = -75
	ProtocolVersionHixie76   = -76
	ProtocolVersionHybi00    = 0
	ProtocolVersionHybi08    = 8
	ProtocolVersionHybi13    = 13
	ProtocolVersionHybi      = ProtocolVersionHybi13
	SupportedProtocolVersion = "13, 8"

	ContinuationFrame = 0
	TextFrame         = 1
	BinaryFrame       = 2
	CloseFrame        = 8
	PingFrame         = 9
	PongFrame         = 10
	UnknownFrame      = 255
)

// WebSocket protocol errors.
type ProtocolError struct {
	ErrorString string
}

func (err *ProtocolError) Error() string { return err.ErrorString }

var (
	ErrBadProtocolVersion   = &ProtocolError{"bad protocol version"}
	ErrBadScheme            = &ProtocolError{"bad scheme"}
	ErrBadStatus            = &ProtocolError{"bad status"}
	ErrBadUpgrade           = &ProtocolError{"missing or bad upgrade"}
	ErrBadWebSocketOrigin   = &ProtocolError{"missing or bad WebSocket-Origin"}
	ErrBadWebSocketLocation = &ProtocolError{"missing or bad WebSocket-Location"}
	ErrBadWebSocketProtocol = &ProtocolError{"missing or bad WebSocket-Protocol"}
	ErrBadWebSocketVersion  = &ProtocolError{"missing or bad WebSocket Version"}
	ErrChallengeResponse    = &ProtocolError{"mismatch challenge/response"}
	ErrBadFrame             = &ProtocolError{"bad frame"}
	ErrBadFrameBoundary     = &ProtocolError{"not on frame boundary"}
	ErrNotWebSocket         = &ProtocolError{"not websocket protocol"}
	ErrBadRequestMethod     = &ProtocolError{"bad method"}
	ErrNotSupported         = &ProtocolError{"not supported"}
)

// Addr is an implementation of net.Addr for WebSocket.
type Addr struct {
	*url.URL
}

// Network returns the network type for a WebSocket, "websocket".
func (addr *Addr) Network() string { return "websocket" }

// Config is a WebSocket configuration
type Config struct {
	// A WebSocket server address.
	Location *url.URL

	// A Websocket client origin.
	Origin *url.URL

	// WebSocket subprotocols.
	Protocol []string

	// WebSocket protocol version.
	Version int

	// TLS config for secure WebSocket (wss).
	TlsConfig *tls.Config

	handshakeData map[string]string
}

// serverHandshaker is an interface to handle WebSocket server side handshake.
type serverHandshaker interface {
	// ReadHandshake reads handshake request message from client.
	// Returns http response code and error if any.
	ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err error)

	// AcceptHandshake accepts the client handshake request and sends
	// handshake response back to client.
	AcceptHandshake(buf *bufio.Writer) (err error)

	// NewServerConn creates a new WebSocket connection.
	NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
}

// frameReader is an interface to read a WebSocket frame.
type frameReader interface {
	// Reader is to read payload of the frame.
	io.Reader

	// PayloadType returns payload type.
	PayloadType() byte

	// HeaderReader returns a reader to read header of the frame.
	HeaderReader() io.Reader

	// TrailerReader returns a reader to read trailer of the frame.
	// If it returns nil, there is no trailer in the frame.
	TrailerReader() io.Reader

	// Len returns total length of the frame, including header and trailer.
	Len() int
}

// frameReaderFactory is an interface to creates new frame reader.
type frameReaderFactory interface {
	NewFrameReader() (r frameReader, err error)
}

// frameWriter is an interface to write a WebSocket frame.
type frameWriter interface {
	// Writer is to write playload of the frame.
	io.WriteCloser
}

// frameWriterFactory is an interface to create new frame writer.
type frameWriterFactory interface {
	NewFrameWriter(payloadType byte) (w frameWriter, err error)
}

type frameHandler interface {
	HandleFrame(frame frameReader) (r frameReader, err error)
	WriteClose(status int) (err error)
}

// Conn represents a WebSocket connection.
type Conn struct {
	config  *Config
	request *http.Request

	buf *bufio.ReadWriter
	rwc io.ReadWriteCloser

	rio sync.Mutex
	frameReaderFactory
	frameReader

	wio sync.Mutex
	frameWriterFactory

	frameHandler
	PayloadType        byte
	defaultCloseStatus int
}

// Read implements the io.Reader interface:
// it reads data of a frame from the WebSocket connection.
// if msg is not large enough for the frame data, it fills the msg and next Read
// will read the rest of the frame data.
// it reads Text frame or Binary frame.
func (ws *Conn) Read(msg []byte) (n int, err error) {
	ws.rio.Lock()
	defer ws.rio.Unlock()
again:
	if ws.frameReader == nil {
		frame, err := ws.frameReaderFactory.NewFrameReader()
		if err != nil {
			return 0, err
		}
		ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
		if err != nil {
			return 0, err
		}
		if ws.frameReader == nil {
			goto again
		}
	}
	n, err = ws.frameReader.Read(msg)
	if err == io.EOF {
		if trailer := ws.frameReader.TrailerReader(); trailer != nil {
			io.Copy(ioutil.Discard, trailer)
		}
		ws.frameReader = nil
		goto again
	}
	return n, err
}

// Write implements the io.Writer interface:
// it writes data as a frame to the WebSocket connection.
func (ws *Conn) Write(msg []byte) (n int, err error) {
	ws.wio.Lock()
	defer ws.wio.Unlock()
	w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
	if err != nil {
		return 0, err
	}
	n, err = w.Write(msg)
	w.Close()
	if err != nil {
		return n, err
	}
	return n, err
}

// Close implements the io.Closer interface.
func (ws *Conn) Close() error {
	err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
	if err != nil {
		return err
	}
	return ws.rwc.Close()
}

func (ws *Conn) IsClientConn() bool { return ws.request == nil }
func (ws *Conn) IsServerConn() bool { return ws.request != nil }

// LocalAddr returns the WebSocket Origin for the connection for client, or
// the WebSocket location for server.
func (ws *Conn) LocalAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Origin}
	}
	return &Addr{ws.config.Location}
}

// RemoteAddr returns the WebSocket location for the connection for client, or
// the Websocket Origin for server.
func (ws *Conn) RemoteAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Location}
	}
	return &Addr{ws.config.Origin}
}

var errSetDeadline = errors.New("websocket: cannot set deadline: not using a net.Conn")

// SetDeadline sets the connection's network read & write deadlines.
func (ws *Conn) SetDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetDeadline(t)
	}
	return errSetDeadline
}

// SetReadDeadline sets the connection's network read deadline.
func (ws *Conn) SetReadDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetReadDeadline(t)
	}
	return errSetDeadline
}

// SetWriteDeadline sets the connection's network write deadline.
func (ws *Conn) SetWriteDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetWriteDeadline(t)
	}
	return errSetDeadline
}

// Config returns the WebSocket config.
func (ws *Conn) Config() *Config { return ws.config }

// Request returns the http request upgraded to the WebSocket.
// It is nil for client side.
func (ws *Conn) Request() *http.Request { return ws.request }

// Codec represents a symmetric pair of functions that implement a codec.
type Codec struct {
	Marshal   func(v interface{}) (data []byte, payloadType byte, err error)
	Unmarshal func(data []byte, payloadType byte, v interface{}) (err error)
}

// Send sends v marshaled by cd.Marshal as single frame to ws.
func (cd Codec) Send(ws *Conn, v interface{}) (err error) {
	data, payloadType, err := cd.Marshal(v)
	if err != nil {
		return err
	}
	ws.wio.Lock()
	defer ws.wio.Unlock()
	w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
	if err != nil {
		return err
	}
	_, err = w.Write(data)
	w.Close()
	return err
}

// Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores in v.
func (cd Codec) Receive(ws *Conn, v interface{}) (err error) {
	ws.rio.Lock()
	defer ws.rio.Unlock()
	if ws.frameReader != nil {
		_, err = io.Copy(ioutil.Discard, ws.frameReader)
		if err != nil {
			return err
		}
		ws.frameReader = nil
	}
again:
	frame, err := ws.frameReaderFactory.NewFrameReader()
	if err != nil {
		return err
	}
	frame, err = ws.frameHandler.HandleFrame(frame)
	if err != nil {
		return err
	}
	if frame == nil {
		goto again
	}
	payloadType := frame.PayloadType()
	data, err := ioutil.ReadAll(frame)
	if err != nil {
		return err
	}
	return cd.Unmarshal(data, payloadType, v)
}

func marshal(v interface{}) (msg []byte, payloadType byte, err error) {
	switch data := v.(type) {
	case string:
		return []byte(data), TextFrame, nil
	case []byte:
		return data, BinaryFrame, nil
	}
	return nil, UnknownFrame, ErrNotSupported
}

func unmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	switch data := v.(type) {
	case *string:
		*data = string(msg)
		return nil
	case *[]byte:
		*data = msg
		return nil
	}
	return ErrNotSupported
}

/*
Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
To send/receive text frame, use string type.
To send/receive binary frame, use []byte type.

Trivial usage:

	import "websocket"

	// receive text frame
	var message string
	websocket.Message.Receive(ws, &message)

	// send text frame
	message = "hello"
	websocket.Message.Send(ws, message)

	// receive binary frame
	var data []byte
	websocket.Message.Receive(ws, &data)

	// send binary frame
	data = []byte{0, 1, 2}
	websocket.Message.Send(ws, data)

*/
var Message = Codec{marshal, unmarshal}

func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err error) {
	msg, err = json.Marshal(v)
	return msg, TextFrame, err
}

func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	return json.Unmarshal(msg, v)
}

/*
JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.

Trival usage:

	import "websocket"

	type T struct {
		Msg string
		Count int
	}

	// receive JSON type T
	var data T
	websocket.JSON.Receive(ws, &data)

	// send JSON type T
	websocket.JSON.Send(ws, data)
*/
var JSON = Codec{jsonMarshal, jsonUnmarshal}
<|endoftext|>


<commit_before>package util

const (
	VERSION = "1.09"
)
<commit_msg>for 1.10<commit_after>package util

const (
	VERSION = "1.10"
)
<|endoftext|>


<commit_before>package main_test

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

	cli "github.com/heroku/cli"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("update.go", func() {
	Describe("UpdateCLI()", func() {
		It("downloads a new CLI", func() {
			os.Remove(cli.UpdateLockPath)
			manifest := cli.GetUpdateManifest("dev")
			Expect(manifest.Channel).To(Equal("dev"))
			dest := filepath.Join("tmp", "newcli")
			cli.DownloadCLI("dev", dest, manifest)
			out, err := exec.Command(filepath.Join(dest, "bin", "heroku"), "version").Output()
			must(err)
			Expect(out).To(HavePrefix("heroku-cli/"))
		})
	})
})
<commit_msg>circle fix<commit_after>package main_test

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

	cli "github.com/heroku/cli"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("update.go", func() {
	Describe("UpdateCLI()", func() {
		It("downloads a new CLI", func() {
			Skip("not working on circle")
			os.Remove(cli.UpdateLockPath)
			manifest := cli.GetUpdateManifest("dev")
			Expect(manifest.Channel).To(Equal("dev"))
			dest := filepath.Join("tmp", "newcli")
			cli.DownloadCLI("dev", dest, manifest)
			out, err := exec.Command(filepath.Join(dest, "bin", "heroku"), "version").Output()
			must(err)
			Expect(out).To(HavePrefix("heroku-cli/"))
		})
	})
})
<|endoftext|>