<commit_before>package proxy

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/hashicorp/consul/agent"
	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/connect"
	"github.com/hashicorp/consul/sdk/testutil"
)

func TestUpstreamResolverFuncFromClient(t *testing.T) {
	tests := []struct {
		name string
		cfg  UpstreamConfig
		want *connect.ConsulResolver
	}{
		{
			name: "service",
			cfg: UpstreamConfig{
				DestinationNamespace: "foo",
				DestinationPartition: "default",
				DestinationName:      "web",
				Datacenter:           "ny1",
				DestinationType:      "service",
			},
			want: &connect.ConsulResolver{
				Namespace:  "foo",
				Partition:  "default",
				Name:       "web",
				Datacenter: "ny1",
				Type:       connect.ConsulResolverTypeService,
			},
		},
		{
			name: "prepared_query",
			cfg: UpstreamConfig{
				DestinationNamespace: "foo",
				DestinationPartition: "default",
				DestinationName:      "web",
				Datacenter:           "ny1",
				DestinationType:      "prepared_query",
			},
			want: &connect.ConsulResolver{
				Namespace:  "foo",
				Name:       "web",
				Partition:  "default",
				Datacenter: "ny1",
				Type:       connect.ConsulResolverTypePreparedQuery,
			},
		},
		{
			name: "unknown behaves like service",
			cfg: UpstreamConfig{
				DestinationNamespace: "foo",
				DestinationPartition: "default",
				DestinationName:      "web",
				Datacenter:           "ny1",
				DestinationType:      "junk",
			},
			want: &connect.ConsulResolver{
				Partition:  "default",
				Namespace:  "foo",
				Name:       "web",
				Datacenter: "ny1",
				Type:       connect.ConsulResolverTypeService,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Client doesn't really matter as long as it's passed through.
			gotFn := UpstreamResolverFuncFromClient(nil)
			got, err := gotFn(tt.cfg)
			require.NoError(t, err)
			require.Equal(t, tt.want, got)
		})
	}
}

func TestAgentConfigWatcherSidecarProxy(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	a := agent.StartTestAgent(t, agent.TestAgent{Name: "agent_smith"})
	defer a.Shutdown()

	client := a.Client()
	agent := client.Agent()

	// Register a local agent service with a sidecar proxy
	reg := &api.AgentServiceRegistration{
		Name: "web",
		Port: 8080,
		Connect: &api.AgentServiceConnect{
			SidecarService: &api.AgentServiceRegistration{
				Proxy: &api.AgentServiceConnectProxyConfig{
					Config: map[string]interface{}{
						"handshake_timeout_ms": 999,
					},
					Upstreams: []api.Upstream{
						{
							DestinationName: "db",
							LocalBindPort:   9191,
						},
					},
				},
			},
		},
	}
	err := agent.ServiceRegister(reg)
	require.NoError(t, err)

	w, err := NewAgentConfigWatcher(client, "web-sidecar-proxy",
		testutil.Logger(t))
	require.NoError(t, err)

	cfg := testGetConfigValTimeout(t, w, 500*time.Millisecond)

	expectCfg := &Config{
		ProxiedServiceName:      "web",
		ProxiedServiceNamespace: "default",
		PublicListener: PublicListenerConfig{
			BindAddress:           "0.0.0.0",
			BindPort:              21000,
			LocalServiceAddress:   "127.0.0.1:8080",
			HandshakeTimeoutMs:    999,
			LocalConnectTimeoutMs: 1000, // from applyDefaults
		},
		Upstreams: []UpstreamConfig{
			{
				DestinationName:      "db",
				DestinationNamespace: "default",
				DestinationPartition: "default",
				DestinationType:      "service",
				LocalBindPort:        9191,
				LocalBindAddress:     "127.0.0.1",
			},
		},
	}
	require.Equal(t, expectCfg, cfg)

	// Now keep watching and update the config.
	reg.Connect.SidecarService.Proxy.Upstreams = append(reg.Connect.SidecarService.Proxy.Upstreams,
		api.Upstream{
			DestinationName:  "cache",
			LocalBindPort:    9292,
			LocalBindAddress: "127.10.10.10",
		})
	reg.Connect.SidecarService.Proxy.Config["local_connect_timeout_ms"] = 444
	require.NoError(t, agent.ServiceRegister(reg))

	cfg = testGetConfigValTimeout(t, w, 2*time.Second)

	expectCfg.Upstreams = append(expectCfg.Upstreams, UpstreamConfig{
		DestinationName:      "cache",
		DestinationNamespace: "default",
		DestinationPartition: "default",
		DestinationType:      "service",
		LocalBindPort:        9292,
		LocalBindAddress:     "127.10.10.10",
	})
	expectCfg.PublicListener.LocalConnectTimeoutMs = 444

	require.Equal(t, expectCfg, cfg)
}

func testGetConfigValTimeout(t *testing.T, w ConfigWatcher,
	timeout time.Duration) *Config {
	t.Helper()
	select {
	case cfg := <-w.Watch():
		return cfg
	case <-time.After(timeout):
		t.Fatalf("timeout after %s waiting for config update", timeout)
		return nil
	}
}
<commit_msg>Add retries and debugging to flaky test<commit_after>package proxy

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/hashicorp/consul/agent"
	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/connect"
	"github.com/hashicorp/consul/sdk/testutil"
	"github.com/hashicorp/consul/sdk/testutil/retry"
)

func TestUpstreamResolverFuncFromClient(t *testing.T) {
	tests := []struct {
		name string
		cfg  UpstreamConfig
		want *connect.ConsulResolver
	}{
		{
			name: "service",
			cfg: UpstreamConfig{
				DestinationNamespace: "foo",
				DestinationPartition: "default",
				DestinationName:      "web",
				Datacenter:           "ny1",
				DestinationType:      "service",
			},
			want: &connect.ConsulResolver{
				Namespace:  "foo",
				Partition:  "default",
				Name:       "web",
				Datacenter: "ny1",
				Type:       connect.ConsulResolverTypeService,
			},
		},
		{
			name: "prepared_query",
			cfg: UpstreamConfig{
				DestinationNamespace: "foo",
				DestinationPartition: "default",
				DestinationName:      "web",
				Datacenter:           "ny1",
				DestinationType:      "prepared_query",
			},
			want: &connect.ConsulResolver{
				Namespace:  "foo",
				Name:       "web",
				Partition:  "default",
				Datacenter: "ny1",
				Type:       connect.ConsulResolverTypePreparedQuery,
			},
		},
		{
			name: "unknown behaves like service",
			cfg: UpstreamConfig{
				DestinationNamespace: "foo",
				DestinationPartition: "default",
				DestinationName:      "web",
				Datacenter:           "ny1",
				DestinationType:      "junk",
			},
			want: &connect.ConsulResolver{
				Partition:  "default",
				Namespace:  "foo",
				Name:       "web",
				Datacenter: "ny1",
				Type:       connect.ConsulResolverTypeService,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Client doesn't really matter as long as it's passed through.
			gotFn := UpstreamResolverFuncFromClient(nil)
			got, err := gotFn(tt.cfg)
			require.NoError(t, err)
			require.Equal(t, tt.want, got)
		})
	}
}

func TestAgentConfigWatcherSidecarProxy(t *testing.T) {
	if testing.Short() {
		t.Skip("too slow for testing.Short")
	}

	a := agent.StartTestAgent(t, agent.TestAgent{Name: "agent_smith"})
	defer a.Shutdown()

	client := a.Client()
	agent := client.Agent()

	// Register a local agent service with a sidecar proxy
	reg := &api.AgentServiceRegistration{
		Name: "web",
		Port: 8080,
		Connect: &api.AgentServiceConnect{
			SidecarService: &api.AgentServiceRegistration{
				Proxy: &api.AgentServiceConnectProxyConfig{
					Config: map[string]interface{}{
						"handshake_timeout_ms": 999,
					},
					Upstreams: []api.Upstream{
						{
							DestinationName: "db",
							LocalBindPort:   9191,
						},
					},
				},
			},
		},
	}
	err := agent.ServiceRegister(reg)
	require.NoError(t, err)

	w, err := NewAgentConfigWatcher(client, "web-sidecar-proxy",
		testutil.Logger(t))
	require.NoError(t, err)

	cfg := testGetConfigValTimeout(t, w, 500*time.Millisecond)

	expectCfg := &Config{
		ProxiedServiceName:      "web",
		ProxiedServiceNamespace: "default",
		PublicListener: PublicListenerConfig{
			BindAddress:           "0.0.0.0",
			BindPort:              21000,
			LocalServiceAddress:   "127.0.0.1:8080",
			HandshakeTimeoutMs:    999,
			LocalConnectTimeoutMs: 1000, // from applyDefaults
		},
		Upstreams: []UpstreamConfig{
			{
				DestinationName:      "db",
				DestinationNamespace: "default",
				DestinationPartition: "default",
				DestinationType:      "service",
				LocalBindPort:        9191,
				LocalBindAddress:     "127.0.0.1",
			},
		},
	}
	require.Equal(t, expectCfg, cfg)

	// Now keep watching and update the config.
	reg.Connect.SidecarService.Proxy.Upstreams = append(reg.Connect.SidecarService.Proxy.Upstreams,
		api.Upstream{
			DestinationName:  "cache",
			LocalBindPort:    9292,
			LocalBindAddress: "127.10.10.10",
		})
	reg.Connect.SidecarService.Proxy.Config["local_connect_timeout_ms"] = 444
	require.NoError(t, agent.ServiceRegister(reg))

	updatedCfg := new(Config)
	*updatedCfg = *expectCfg
	updatedCfg.Upstreams = append(updatedCfg.Upstreams, UpstreamConfig{
		DestinationName:      "cache",
		DestinationNamespace: "default",
		DestinationPartition: "default",
		DestinationType:      "service",
		LocalBindPort:        9292,
		LocalBindAddress:     "127.10.10.10",
	})
	updatedCfg.PublicListener.LocalConnectTimeoutMs = 444

	retry.Run(t, func(r *retry.R) {
		cfg := testGetConfigValTimeout(r, w, 500*time.Millisecond)
		// TODO: These are debug logs to show the diffs against updatedCfg and expectCfg.
		// Once we figure out what event makes this test flake, we should adjust this test to be deterministic.
		if !assert.Equal(r, updatedCfg, cfg, "expected config from watcher to match updated") {
			assert.Equal(r, expectCfg, cfg, "config does not match original or updated config; something else must have fired watch")
			r.FailNow()
		}
	})
}

type testingT interface {
	Helper()
	Fatalf(format string, args ...interface{})
}

func testGetConfigValTimeout(t testingT, w ConfigWatcher,
	timeout time.Duration) *Config {
	t.Helper()
	select {
	case cfg := <-w.Watch():
		return cfg
	case <-time.After(timeout):
		t.Fatalf("timeout after %s waiting for config update", timeout)
		return nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/thoj/go-ircevent"
	"os/exec"
)

type Command struct {
	Name     string
	Function func(bot *MusicBot, event *irc.Event, parameters []string)
}

func (c *Command) execute(bot *MusicBot, event *irc.Event, parameters []string) {
	c.Function(bot, event, parameters)
}

var HelpCommand = Command{
	Name: "Help",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		channel := event.Arguments[0]
		message := "Available commands: "
		for commandName, _ := range bot.Commands {
			message += " " + commandName
		}

		event.Connection.Privmsg(channel, message)
	},
}

var WhitelistCommand = Command{
	Name: "Whitelist",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		channel := event.Arguments[0]
		realname := event.User
		if len(parameters) < 1 {
			event.Connection.Privmsg(channel, "!music whitelist <show|add|remove> [user]")
			return
		}

		subcommand := parameters[0]
		switch subcommand {
		case "show":
			{
				message := "Current whitelist: "
				for _, name := range bot.Whitelist {
					message += " " + name
				}

				event.Connection.Privmsg(channel, message)
			}
		case "add":
			{
				if len(parameters) < 2 {
					event.Connection.Privmsg(channel, "!music whitelist add [user]")
					return
				}
				user := parameters[1]
				if realname == bot.Master {
					if isWhitelisted, _ := bot.isUserWhitelisted(user); !isWhitelisted {
						bot.Whitelist = append(bot.Whitelist, user)
						event.Connection.Privmsg(channel, "User added to whitelist")
						writeWhitelist(bot.Whitelist)
					}
				}
			}
		case "remove":
			{
				if len(parameters) < 2 {
					event.Connection.Privmsg(channel, "!music whitelist remove [user]")
					return
				}
				user := parameters[1]
				if realname == bot.Master {
					if isWhitelisted, index := bot.isUserWhitelisted(user); isWhitelisted {
						bot.Whitelist = append(bot.Whitelist[:index], bot.Whitelist[index+1:]...)
						event.Connection.Privmsg(channel, "User removed from whitelist")
						writeWhitelist(bot.Whitelist)
					}
				}
			}
		}
	},
}

var NextCommand = Command{
	Name: "Next",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		bot.MusicPlayer.Next()
	},
}

var PlayCommand = Command{
	Name: "Play",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		bot.MusicPlayer.Pause()
	},
}

var PauseCommand = Command{
	Name: "Pause",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		bot.MusicPlayer.Pause()
	},
}

var CurrentCommand = Command{
	Name: "Current",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		channel := event.Arguments[0]
		url := bot.MusicPlayer.CurrentSong.GetUrl()
		message := fmt.Sprintf("Current song: %s", url)
		event.Connection.Privmsg(channel, message)
	},
}

var OpenCommand = Command{
	Name: "Open",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		if len(parameters) < 1 {
			channel := event.Arguments[0]
			event.Connection.Privmsg(channel, "!music open <youtube link>")
			return
		}

		url := parameters[0]
		fmt.Println(url)
		bot.MusicPlayer.AddSong(url)
	},
}

var SearchCommand = Command{
	Name: "search",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		channel := event.Arguments[0]
		if len(parameters) < 1 {
			event.Connection.Privmsg(channel, "!music search <search term>")
			return
		}
		event.Connection.Privmsg(channel, "Not implemented")
	},
}

var VolUpCommand = Command{
	Name: "vol++",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		cmd := exec.Command("amixer", "-D", "pulse", "sset", "Master", "10%+")
		cmd.Run()
	},
}

var VolDownCommand = Command{
	Name: "vol--",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		cmd := exec.Command("amixer", "-D", "pulse", "sset", "Master", "10%-")
		cmd.Run()
	},
}

var VolCommand = Command{
	Name: "vol",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		if len(parameters) < 1 {
			channel := event.Arguments[0]
			event.Connection.Privmsg(channel, "!music vol <volume>")
			return
		}
		cmd := exec.Command("amixer", "-D", "pulse", "sset", "Master", parameters[0]+"%")
		cmd.Run()
	},
}
<commit_msg>Removed redundant underscore<commit_after>package main

import (
	"fmt"
	"github.com/thoj/go-ircevent"
	"os/exec"
)

type Command struct {
	Name     string
	Function func(bot *MusicBot, event *irc.Event, parameters []string)
}

func (c *Command) execute(bot *MusicBot, event *irc.Event, parameters []string) {
	c.Function(bot, event, parameters)
}

var HelpCommand = Command{
	Name: "Help",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		channel := event.Arguments[0]
		message := "Available commands: "
		for commandName := range bot.Commands {
			message += " " + commandName
		}

		event.Connection.Privmsg(channel, message)
	},
}

var WhitelistCommand = Command{
	Name: "Whitelist",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		channel := event.Arguments[0]
		realname := event.User
		if len(parameters) < 1 {
			event.Connection.Privmsg(channel, "!music whitelist <show|add|remove> [user]")
			return
		}

		subcommand := parameters[0]
		switch subcommand {
		case "show":
			{
				message := "Current whitelist: "
				for _, name := range bot.Whitelist {
					message += " " + name
				}

				event.Connection.Privmsg(channel, message)
			}
		case "add":
			{
				if len(parameters) < 2 {
					event.Connection.Privmsg(channel, "!music whitelist add [user]")
					return
				}
				user := parameters[1]
				if realname == bot.Master {
					if isWhitelisted, _ := bot.isUserWhitelisted(user); !isWhitelisted {
						bot.Whitelist = append(bot.Whitelist, user)
						event.Connection.Privmsg(channel, "User added to whitelist")
						writeWhitelist(bot.Whitelist)
					}
				}
			}
		case "remove":
			{
				if len(parameters) < 2 {
					event.Connection.Privmsg(channel, "!music whitelist remove [user]")
					return
				}
				user := parameters[1]
				if realname == bot.Master {
					if isWhitelisted, index := bot.isUserWhitelisted(user); isWhitelisted {
						bot.Whitelist = append(bot.Whitelist[:index], bot.Whitelist[index+1:]...)
						event.Connection.Privmsg(channel, "User removed from whitelist")
						writeWhitelist(bot.Whitelist)
					}
				}
			}
		}
	},
}

var NextCommand = Command{
	Name: "Next",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		bot.MusicPlayer.Next()
	},
}

var PlayCommand = Command{
	Name: "Play",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		bot.MusicPlayer.Pause()
	},
}

var PauseCommand = Command{
	Name: "Pause",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		bot.MusicPlayer.Pause()
	},
}

var CurrentCommand = Command{
	Name: "Current",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		channel := event.Arguments[0]
		url := bot.MusicPlayer.CurrentSong.GetUrl()
		message := fmt.Sprintf("Current song: %s", url)
		event.Connection.Privmsg(channel, message)
	},
}

var OpenCommand = Command{
	Name: "Open",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		if len(parameters) < 1 {
			channel := event.Arguments[0]
			event.Connection.Privmsg(channel, "!music open <youtube link>")
			return
		}

		url := parameters[0]
		fmt.Println(url)
		bot.MusicPlayer.AddSong(url)
	},
}

var SearchCommand = Command{
	Name: "search",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		channel := event.Arguments[0]
		if len(parameters) < 1 {
			event.Connection.Privmsg(channel, "!music search <search term>")
			return
		}
		event.Connection.Privmsg(channel, "Not implemented")
	},
}

var VolUpCommand = Command{
	Name: "vol++",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		cmd := exec.Command("amixer", "-D", "pulse", "sset", "Master", "10%+")
		cmd.Run()
	},
}

var VolDownCommand = Command{
	Name: "vol--",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		cmd := exec.Command("amixer", "-D", "pulse", "sset", "Master", "10%-")
		cmd.Run()
	},
}

var VolCommand = Command{
	Name: "vol",
	Function: func(bot *MusicBot, event *irc.Event, parameters []string) {
		if len(parameters) < 1 {
			channel := event.Arguments[0]
			event.Connection.Privmsg(channel, "!music vol <volume>")
			return
		}
		cmd := exec.Command("amixer", "-D", "pulse", "sset", "Master", parameters[0]+"%")
		cmd.Run()
	},
}
<|endoftext|>


<commit_before>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func(string) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)

	handlers     map[string]func(*Server, *Client, string) string
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID(r *http.Request) string {
	hash := fmt.Sprintf("%s%d", r.RemoteAddr, time.Now().UnixNano())
	return fmt.Sprintf("%x", sha1.Sum([]byte(hash)))
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       genID(r),
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

func (c *Client) Close() error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.Close()
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := struct {
		Id   int         `json:"id"`
		Data interface{} `json:"data"`
	}{-1, v}

	return c.conn.WriteJSON(js)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func(*Server, *Client, string) string)),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	go s.listener()
	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		if r == room {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.Lock()

	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		s.roomsLock.Unlock()
		return
	}

	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	s.roomsLock.Unlock()

	index = -1

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	if index == -1 {
		return
	}

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
}

//Send all clients in room room data with type messageType
func (s *Server) Broadcast(room string, data string) {
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			c.Emit(data)
		}(client)
	}
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		go func(c *Client) {
			c.EmitJSON(v)
		}(client)
	}
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.roomsLock.Lock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}
		if index == -1 {
			s.roomsLock.Unlock()
			return
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
	}
	s.roomsLock.Unlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	var rooms []string
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}

func (c *Client) listener(s *Server) {
	for {
		mtype, data, err := c.conn.ReadMessage()
		if err != nil {
			c.cleanup(s)
			return
		}

		var js struct {
			Id   string
			Data json.RawMessage
		}
		err = json.Unmarshal(data, &js)

		if err != nil || mtype != ws.TextMessage {
			log.Println(err)
			continue
		}

		callName := s.Extractor(string(js.Data))

		s.handlersLock.RLock()
		f, ok := s.handlers[callName]
		s.handlersLock.RUnlock()

		if !ok {
			continue
		}

		go func() {
			rtrn := f(s, c, string(js.Data))
			reply := struct {
				Id   string `json:"id"`
				Data string `json:"data,string"`
			}{js.Id, rtrn}

			bytes, _ := json.Marshal(reply)
			c.Emit(string(bytes))
		}()
	}
}

func (s *Server) listener() {
	for {
		select {
		case c := <-s.newClient:
			go c.listener(s)
		default:
			log.Fatal("newClient channel closed")
		}

	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f func(*Server, *Client, string) string) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<commit_msg>Don't exit on default case.<commit_after>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func(string) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)

	handlers     map[string]func(*Server, *Client, string) string
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID(r *http.Request) string {
	hash := fmt.Sprintf("%s%d", r.RemoteAddr, time.Now().UnixNano())
	return fmt.Sprintf("%x", sha1.Sum([]byte(hash)))
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       genID(r),
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

func (c *Client) Close() error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.Close()
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := struct {
		Id   int         `json:"id"`
		Data interface{} `json:"data"`
	}{-1, v}

	return c.conn.WriteJSON(js)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func(*Server, *Client, string) string)),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	go s.listener()
	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		if r == room {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.Lock()

	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		s.roomsLock.Unlock()
		return
	}

	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	s.roomsLock.Unlock()

	index = -1

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	if index == -1 {
		return
	}

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
}

//Send all clients in room room data with type messageType
func (s *Server) Broadcast(room string, data string) {
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			c.Emit(data)
		}(client)
	}
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		go func(c *Client) {
			c.EmitJSON(v)
		}(client)
	}
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.roomsLock.Lock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}
		if index == -1 {
			s.roomsLock.Unlock()
			return
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
	}
	s.roomsLock.Unlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	var rooms []string
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}

func (c *Client) listener(s *Server) {
	for {
		mtype, data, err := c.conn.ReadMessage()
		if err != nil {
			c.cleanup(s)
			return
		}

		var js struct {
			Id   string
			Data json.RawMessage
		}
		err = json.Unmarshal(data, &js)

		if err != nil || mtype != ws.TextMessage {
			log.Println(err)
			continue
		}

		callName := s.Extractor(string(js.Data))

		s.handlersLock.RLock()
		f, ok := s.handlers[callName]
		s.handlersLock.RUnlock()

		if !ok {
			continue
		}

		go func() {
			rtrn := f(s, c, string(js.Data))
			reply := struct {
				Id   string `json:"id"`
				Data string `json:"data,string"`
			}{js.Id, rtrn}

			bytes, _ := json.Marshal(reply)
			c.Emit(string(bytes))
		}()
	}
}

func (s *Server) listener() {
	for {
		select {
		case c := <-s.newClient:
			go c.listener(s)
		default:
		}

	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f func(*Server, *Client, string) string) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<|endoftext|>


<commit_before>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"log"
	"net/http"
	"sync"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func(string) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)
	//Called when no event handler for a specific event exists
	DefaultHandler func(*Server, *Client, []byte) []byte

	handlers     map[string]func(*Server, *Client, []byte) []byte
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID() string {
	bytes := make([]byte, 32)
	rand.Read(bytes)

	return base64.URLEncoding.EncodeToString(bytes)
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClientWithID(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request, id string) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       id,
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	return s.NewClientWithID(upgrader, w, r, genID())
}

func (c *Client) Close() error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.Close()
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := struct {
		Id   int         `json:"id"`
		Data interface{} `json:"data"`
	}{-1, v}

	return c.conn.WriteJSON(js)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func(*Server, *Client, []byte) []byte)),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	go s.listener()
	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		if r == room {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.Lock()

	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		s.roomsLock.Unlock()
		return
	}

	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	s.roomsLock.Unlock()

	index = -1

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	if index == -1 {
		return
	}

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
}

//Send all clients in room room data with type messageType
func (s *Server) Broadcast(room string, data string) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			c.Emit(data)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		go func(c *Client) {
			c.EmitJSON(v)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		s.roomsLock.Lock()
		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
		s.roomsLock.Unlock()
	}
	s.joinedRoomsLock.RUnlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	rooms := make([]string, len(s.joinedRooms[id]))
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}

func (c *Client) listener(s *Server) {
	for {
		mtype, data, err := c.conn.ReadMessage()
		if err != nil {
			c.cleanup(s)
			return
		}

		var js struct {
			Id   string
			Data json.RawMessage
		}
		err = json.Unmarshal(data, &js)

		if err != nil || mtype != ws.TextMessage {
			log.Println(err)
			continue
		}

		callName := s.Extractor(string(js.Data))

		s.handlersLock.RLock()
		f, ok := s.handlers[callName]
		s.handlersLock.RUnlock()

		if !ok {
			if s.DefaultHandler != nil {
				f = s.DefaultHandler
				goto call
			}
			continue
		}
	call:
		go func() {
			rtrn := f(s, c, js.Data)
			reply := struct {
				Id   string `json:"id"`
				Data string `json:"data,string"`
			}{js.Id, string(rtrn)}

			bytes, _ := json.Marshal(reply)
			c.Emit(string(bytes))
		}()
	}
}

func (s *Server) listener() {
	for {
		c := <-s.newClient
		go c.listener(s)
	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f func(*Server, *Client, []byte) []byte) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<commit_msg>Use sync.Pool<commit_after>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"log"
	"net/http"
	"sync"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func(string) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)
	//Called when no event handler for a specific event exists
	DefaultHandler func(*Server, *Client, []byte) []byte

	handlers     map[string]func(*Server, *Client, []byte) []byte
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID() string {
	bytes := make([]byte, 32)
	rand.Read(bytes)

	return base64.URLEncoding.EncodeToString(bytes)
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClientWithID(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request, id string) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       id,
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	return s.NewClientWithID(upgrader, w, r, genID())
}

func (c *Client) Close() error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.Close()
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := struct {
		Id   int         `json:"id"`
		Data interface{} `json:"data"`
	}{-1, v}

	return c.conn.WriteJSON(js)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func(*Server, *Client, []byte) []byte)),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	go s.listener()
	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		if r == room {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.Lock()

	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		s.roomsLock.Unlock()
		return
	}

	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	s.roomsLock.Unlock()

	index = -1

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	if index == -1 {
		return
	}

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
}

//Send all clients in room room data
func (s *Server) Broadcast(room string, data string) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			c.Emit(data)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		go func(c *Client) {
			c.EmitJSON(v)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		s.roomsLock.Lock()
		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
		s.roomsLock.Unlock()
	}
	s.joinedRoomsLock.RUnlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	rooms := make([]string, len(s.joinedRooms[id]))
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}

type request struct {
	Id   string
	Data json.RawMessage
}

type reply struct {
	Id   string `json:"id"`
	Data string `json:"data,string"`
}

var (
	reqPool   = &sync.Pool{New: func() interface{} { return request{} }}
	replyPool = &sync.Pool{New: func() interface{} { return reply{} }}
)

func (c *Client) listener(s *Server) {
	for {
		mtype, data, err := c.conn.ReadMessage()
		if err != nil {
			c.cleanup(s)
			return
		}

		js := reqPool.Get().(request)
		err = json.Unmarshal(data, &js)

		if err != nil || mtype != ws.TextMessage {
			log.Println(err)
			continue
		}

		callName := s.Extractor(string(js.Data))

		s.handlersLock.RLock()
		f, ok := s.handlers[callName]
		s.handlersLock.RUnlock()

		if !ok {
			if s.DefaultHandler != nil {
				f = s.DefaultHandler
				goto call
			}
			continue
		}
	call:
		go func() {
			rtrn := f(s, c, js.Data)

			replyJs := replyPool.Get().(reply)
			replyJs.Id = js.Id
			replyJs.Data = string(rtrn)

			bytes, _ := json.Marshal(replyJs)
			c.Emit(string(bytes))

			reqPool.Put(js)
			replyPool.Put(replyJs)
		}()
	}
}

func (s *Server) listener() {
	for {
		c := <-s.newClient
		go c.listener(s)
	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f func(*Server, *Client, []byte) []byte) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<|endoftext|>


<commit_before>// Copyright 2015 Andrew E. Bruno. All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package nwalgo

import (
	"fmt"
)

var UP = 1
var LEFT = 2
var NW = 3
var NONE = 4

func Align(a, b string, match, mismatch, gap int) (string, string, int) {

	alen := len(a) + 1
	blen := len(b) + 1

	f := make([][]int, alen)
	pointer := make([][]int, alen)
	for i := range f {
		f[i] = make([]int, blen)
		pointer[i] = make([]int, blen)
	}

	for i := 1; i < alen; i++ {
		f[i][0] = gap * i
		pointer[i][0] = UP
	}
	for j := 1; j < blen; j++ {
		f[0][j] = gap * j
		pointer[0][j] = LEFT
	}

	pointer[0][0] = NONE
	for i := 1; i < alen; i++ {
		for j := 1; j < blen; j++ {
			match_mismatch := mismatch
			if a[i-1] == b[j-1] {
				match_mismatch = match
			}

			max := f[i-1][j-1] + match_mismatch
			hgap := f[i-1][j] + gap
			vgap := f[i][j-1] + gap
			if hgap > max {
				max = hgap
			}
			if vgap > max {
				max = vgap
			}

			p := NW
			if max == hgap {
				p = UP
			} else if max == vgap {
				p = LEFT
			}

			pointer[i][j] = p
			f[i][j] = max
		}
	}

	i := alen - 1
	j := blen - 1
	aln1 := ""
	aln2 := ""
	score := f[i][j]

	for p := pointer[i][j]; p != NONE; p = pointer[i][j] {
		if p == NW {
			aln1 = fmt.Sprintf("%c%s", a[i-1], aln1)
			aln2 = fmt.Sprintf("%c%s", b[j-1], aln2)
			i--
			j--
		} else if p == UP {
			aln1 = fmt.Sprintf("%c%s", a[i-1], aln1)
			aln2 = fmt.Sprintf("%s%s", "-", aln2)
			i--
		} else if p == LEFT {
			aln1 = fmt.Sprintf("%s%s", "-", aln1)
			aln2 = fmt.Sprintf("%c%s", b[j-1], aln2)
			j--
		}
	}

	return aln1, aln2, score
}
<commit_msg>Use const for constants.<commit_after>// Copyright 2015 Andrew E. Bruno. All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package nwalgo

import (
	"fmt"
)

const (
	Up   = 1
	Left = 2
	NW   = 3
	None = 4
)

func Align(a, b string, match, mismatch, gap int) (string, string, int) {

	alen := len(a) + 1
	blen := len(b) + 1

	f := make([][]int, alen)
	pointer := make([][]int, alen)
	for i := range f {
		f[i] = make([]int, blen)
		pointer[i] = make([]int, blen)
	}

	for i := 1; i < alen; i++ {
		f[i][0] = gap * i
		pointer[i][0] = Up
	}
	for j := 1; j < blen; j++ {
		f[0][j] = gap * j
		pointer[0][j] = Left
	}

	pointer[0][0] = None
	for i := 1; i < alen; i++ {
		for j := 1; j < blen; j++ {
			match_mismatch := mismatch
			if a[i-1] == b[j-1] {
				match_mismatch = match
			}

			max := f[i-1][j-1] + match_mismatch
			hgap := f[i-1][j] + gap
			vgap := f[i][j-1] + gap
			if hgap > max {
				max = hgap
			}
			if vgap > max {
				max = vgap
			}

			p := NW
			if max == hgap {
				p = Up
			} else if max == vgap {
				p = Left
			}

			pointer[i][j] = p
			f[i][j] = max
		}
	}

	i := alen - 1
	j := blen - 1
	aln1 := ""
	aln2 := ""
	score := f[i][j]

	for p := pointer[i][j]; p != None; p = pointer[i][j] {
		if p == NW {
			aln1 = fmt.Sprintf("%c%s", a[i-1], aln1)
			aln2 = fmt.Sprintf("%c%s", b[j-1], aln2)
			i--
			j--
		} else if p == Up {
			aln1 = fmt.Sprintf("%c%s", a[i-1], aln1)
			aln2 = fmt.Sprintf("%s%s", "-", aln2)
			i--
		} else if p == Left {
			aln1 = fmt.Sprintf("%s%s", "-", aln1)
			aln2 = fmt.Sprintf("%c%s", b[j-1], aln2)
			j--
		}
	}

	return aln1, aln2, score
}
<|endoftext|>


<commit_before>package colorize

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
)

/*
\033[0m 关闭所有属性
\033[1m 设置高亮度
\033[4m 下划线
\033[5m 闪烁
\033[7m 反显
\033[8m 消隐
\033[30m 至 \33[37m 设置前景色
\033[40m 至 \33[47m 设置背景色
\033[nA 光标上移n行
\033[nB 光标下移n行
\033[nC 光标右移n行
\033[nD 光标左移n行
\033[y;xH设置光标位置
\033[2J 清屏
\033[K 清除从光标到行尾的内容
\033[s 保存光标位置
\033[u 恢复光标位置
\033[?25l 隐藏光标
\033[?25h 显示光标
*/

type console_t int
type color_t int

const (
	BLACK color_t = iota
	RED
	GREEN
	YELLOW
	BLUE
	MAGENTA
	CYAN
	WHITE
	DEFAULT color_t = 9
	FORE    color_t = 30
	BACK    color_t = 40
)

const (
	Blod       console_t = 1
	Dim        console_t = 2
	Underlined console_t = 4
	Blinking   console_t = 5
	Reverse    console_t = 7
	Hidden     console_t = 8
)

const prefix = "\033["
const reset console_t = 20
const clear console_t = 0

type Writer struct {
	Attrs      []console_t
	ResetAttrs []console_t
	Fore       color_t
	Back       color_t
	AutoClear  bool
	Escape     bool
	IsHtml     bool
	Writer     io.Writer
}

func NewWriter(w interface{}) *Writer {
	wr := &Writer{
		Fore:   DEFAULT,
		Back:   DEFAULT,
		Writer: os.Stdout,
	}
	if w, ok := w.(io.Writer); ok {
		wr.Writer = w
	}

	if _, ok := w.(*os.File); ok {
		wr.Escape = true
	}
	return wr
}

func (w *Writer) AddAttr(s console_t) {
	w.Attrs = append(w.Attrs, s)
}

func (w *Writer) AddRestAttr(s console_t) {
	w.ResetAttrs = append(w.ResetAttrs, s)
}

func (w *Writer) ClearAttrs() {
	w.Attrs = w.Attrs[0:0]
	w.ResetAttrs = w.ResetAttrs[0:0]
}

func (w *Writer) formatFore() string {
	return fmt.Sprintf("%d", FORE+w.Fore)
}

func (w *Writer) formatBack() string {
	return fmt.Sprintf("%d", BACK+w.Back)
}

func (w *Writer) formatAttrs() string {
	format := ""
	for _, attr := range w.Attrs {
		format += fmt.Sprintf("%d,", attr)
	}

	return strings.TrimRight(format, ",")
}

func (w *Writer) left() string {
	left := prefix + w.formatFore() + ";" + w.formatBack() + ";" + w.formatAttrs()
	left = strings.TrimRight(left, ";") + "m"
	return left
}

func (w *Writer) right() string {
	return prefix + fmt.Sprintf("%d", clear) + "m"
}

func (w *Writer) WriteString(msg string) (int, error) {
	return w.Write([]byte(msg))
}

func (w *Writer) Write(msg []byte) (int, error) {
	if w.Escape {
		return w.Writer.Write(msg)
	}

	buf := w.left()
	_, err := w.Writer.Write([]byte(buf))
	if err != nil {
		return 0, errors.New("format error")
	}

	n, err := w.Writer.Write(msg)
	if err != nil {
		return n, err
	}

	buf = w.right()
	_, err = w.Writer.Write([]byte(buf))
	if err != nil {
		return n, errors.New("format error")
	}
	return n, nil
}

func (w *Writer) converToHtml(msg string) string {
	return ""
}
<commit_msg>fix format<commit_after>package colorize

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
)

/*
\033[0m 关闭所有属性
\033[1m 设置高亮度
\033[4m 下划线
\033[5m 闪烁
\033[7m 反显
\033[8m 消隐
\033[30m 至 \33[37m 设置前景色
\033[40m 至 \33[47m 设置背景色
\033[nA 光标上移n行
\033[nB 光标下移n行
\033[nC 光标右移n行
\033[nD 光标左移n行
\033[y;xH设置光标位置
\033[2J 清屏
\033[K 清除从光标到行尾的内容
\033[s 保存光标位置
\033[u 恢复光标位置
\033[?25l 隐藏光标
\033[?25h 显示光标
*/

type console_t int
type color_t int

const (
	BLACK color_t = iota
	RED
	GREEN
	YELLOW
	BLUE
	MAGENTA
	CYAN
	WHITE
	DEFAULT color_t = 9
	FORE    color_t = 30
	BACK    color_t = 40
)

const (
	Blod       console_t = 1
	Dim        console_t = 2
	Underlined console_t = 4
	Blinking   console_t = 5
	Reverse    console_t = 7
	Hidden     console_t = 8
)

const prefix = "\033["
const reset console_t = 20
const clear console_t = 0

type Writer struct {
	Attrs      []console_t
	ResetAttrs []console_t
	Fore       color_t
	Back       color_t
	AutoClear  bool
	Escape     bool
	IsHtml     bool
	Writer     io.Writer
}

func NewWriter(w interface{}) *Writer {
	wr := &Writer{
		Fore:   DEFAULT,
		Back:   DEFAULT,
		Writer: os.Stdout,
	}
	if w, ok := w.(io.Writer); ok {
		wr.Writer = w
	}

	if _, ok := w.(*os.File); ok {
		wr.Escape = true
	}
	return wr
}

func (w *Writer) AddAttr(s console_t) {
	w.Attrs = append(w.Attrs, s)
}

func (w *Writer) AddRestAttr(s console_t) {
	w.ResetAttrs = append(w.ResetAttrs, s)
}

func (w *Writer) ClearAttrs() {
	w.Attrs = w.Attrs[0:0]
	w.ResetAttrs = w.ResetAttrs[0:0]
}

func (w *Writer) formatFore() string {
	return fmt.Sprintf("%d", FORE+w.Fore)
}

func (w *Writer) formatBack() string {
	return fmt.Sprintf("%d", BACK+w.Back)
}

func (w *Writer) formatAttrs() string {
	format := ""
	for _, attr := range w.Attrs {
		format += fmt.Sprintf("%d;", attr)
	}
	return strings.TrimRight(format, ";")
}

func (w *Writer) left() string {
	left := prefix + w.formatFore() + ";" + w.formatBack() + ";" + w.formatAttrs()
	left = strings.TrimRight(left, ";") + "m"
	return left
}

func (w *Writer) right() string {
	return prefix + fmt.Sprintf("%d", clear) + "m"
}

func (w *Writer) WriteString(msg string) (int, error) {
	return w.Write([]byte(msg))
}

func (w *Writer) Write(msg []byte) (int, error) {
	if w.Escape {
		return w.Writer.Write(msg)
	}

	buf := w.left()
	_, err := w.Writer.Write([]byte(buf))
	if err != nil {
		return 0, errors.New("format error")
	}

	n, err := w.Writer.Write(msg)
	if err != nil {
		return n, err
	}

	buf = w.right()
	_, err = w.Writer.Write([]byte(buf))
	if err != nil {
		return n, errors.New("format error")
	}
	return n, nil
}

func (w *Writer) converToHtml(msg string) string {
	return ""
}
<|endoftext|>


<commit_before>/*
   Copyright 2012 John Asmuth

   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 xgb

import (
	"code.google.com/p/jamslam-x-go-binding/xgb"
	"image/draw"
	"image"
	"fmt"
	"github.com/BurntSushi/xgbutil"
	"github.com/BurntSushi/xgbutil/xgraphics"
	"github.com/BurntSushi/xgbutil/xwindow"
	"github.com/BurntSushi/xgbutil/icccm"
	"github.com/BurntSushi/xgbutil/ewmh"
	"github.com/BurntSushi/xgbutil/keybind"
)

const AllEventsMask =
	xgb.EventMaskKeyPress |
	xgb.EventMaskKeyRelease |
	xgb.EventMaskButtonPress |
	xgb.EventMaskButtonRelease |
    xgb.EventMaskEnterWindow |
    xgb.EventMaskLeaveWindow |
    xgb.EventMaskPointerMotion |
    xgb.EventMaskStructureNotify

type Window struct {
	id xgb.Id
	xu *xgbutil.XUtil
	conn *xgb.Conn
	buffer draw.Image
	width, height int
	closed bool

	events chan interface{}
}

func NewWindow(width, height int) (w *Window, err error) {

	w = new(Window)
	w.width, w.height = width, height
	w.buffer = image.NewRGBA(image.Rectangle{image.Point{0, 0}, image.Point{width, height}})

	w.xu, err = xgbutil.Dial(":0.0")
	if err != nil {
		return
	}

	w.conn = w.xu.Conn()
	screen := w.xu.Screen()

	w.id = w.conn.NewId()
	w.conn.CreateWindow(xgb.WindowClassCopyFromParent, w.id, screen.Root, 600, 500, uint16(width), uint16(height), 0, xgb.WindowClassInputOutput, screen.RootVisual, 0, []uint32{})

	xwindow.Listen(w.xu, w.id, AllEventsMask)

	keyMap, modMap := keybind.MapsGet(w.xu)
	w.xu.KeyMapSet(keyMap)
	w.xu.ModMapSet(modMap)

	w.events = make(chan interface{})

	w.SetIcon(Gordon)
	w.SetIconName("Go")

	go w.handleEvents()

	return
}

func (w *Window) SetTitle(title string) {
	if w.closed {
		return
	}
	err := ewmh.WmNameSet(w.xu, w.id, title)
	if err != nil {
		// TODO: log
	}
	return
}

func (w *Window) SetSize(width, height int) {
	if w.closed {
		return
	}
	err := xwindow.Resize(w.xu, w.id, width, height)
	if err != nil {
		// TODO: log
	}
	w.width, w.height = width, height
	return
}

func (w *Window) Size() (width, height int)  {
	if w.closed {
		return
	}
	width, height = w.width, w.height
	return
}

func (w *Window) Show() {
	if w.closed {
		return
	}
	w.conn.MapWindow(w.id)
	if true {
		err := icccm.WmProtocolsSet(w.xu, w.id, []string{"WM_DELETE_WINDOW"})
		if err != nil {
			fmt.Println(err)
		}
	}
}

func (w *Window) Screen() (im draw.Image) {
	if w.closed {
		return
	}
	im = w.buffer
	return
}

func (w *Window) FlushImage() {
	if w.closed {
		return
	}
	xgraphics.PaintImg(w.xu, w.id, w.buffer)
}

func (w *Window) Close() (err error) {
	if w.closed {
		return
	}
	w.conn.DestroyWindow(w.id)
	w.closed = true
	return
}<commit_msg>don't force DISPLAY<commit_after>/*
   Copyright 2012 John Asmuth

   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 xgb

import (
	"code.google.com/p/jamslam-x-go-binding/xgb"
	"image/draw"
	"image"
	"fmt"
	"github.com/BurntSushi/xgbutil"
	"github.com/BurntSushi/xgbutil/xgraphics"
	"github.com/BurntSushi/xgbutil/xwindow"
	"github.com/BurntSushi/xgbutil/icccm"
	"github.com/BurntSushi/xgbutil/ewmh"
	"github.com/BurntSushi/xgbutil/keybind"
)

const AllEventsMask =
	xgb.EventMaskKeyPress |
	xgb.EventMaskKeyRelease |
	xgb.EventMaskButtonPress |
	xgb.EventMaskButtonRelease |
    xgb.EventMaskEnterWindow |
    xgb.EventMaskLeaveWindow |
    xgb.EventMaskPointerMotion |
    xgb.EventMaskStructureNotify

type Window struct {
	id xgb.Id
	xu *xgbutil.XUtil
	conn *xgb.Conn
	buffer draw.Image
	width, height int
	closed bool

	events chan interface{}
}

func NewWindow(width, height int) (w *Window, err error) {

	w = new(Window)
	w.width, w.height = width, height
	w.buffer = image.NewRGBA(image.Rectangle{image.Point{0, 0}, image.Point{width, height}})

	w.xu, err = xgbutil.Dial("")
	if err != nil {
		return
	}

	w.conn = w.xu.Conn()
	screen := w.xu.Screen()

	w.id = w.conn.NewId()
	w.conn.CreateWindow(xgb.WindowClassCopyFromParent, w.id, screen.Root, 600, 500, uint16(width), uint16(height), 0, xgb.WindowClassInputOutput, screen.RootVisual, 0, []uint32{})

	xwindow.Listen(w.xu, w.id, AllEventsMask)

	keyMap, modMap := keybind.MapsGet(w.xu)
	w.xu.KeyMapSet(keyMap)
	w.xu.ModMapSet(modMap)

	w.events = make(chan interface{})

	w.SetIcon(Gordon)
	w.SetIconName("Go")

	go w.handleEvents()

	return
}

func (w *Window) SetTitle(title string) {
	if w.closed {
		return
	}
	err := ewmh.WmNameSet(w.xu, w.id, title)
	if err != nil {
		// TODO: log
	}
	return
}

func (w *Window) SetSize(width, height int) {
	if w.closed {
		return
	}
	err := xwindow.Resize(w.xu, w.id, width, height)
	if err != nil {
		// TODO: log
	}
	w.width, w.height = width, height
	return
}

func (w *Window) Size() (width, height int)  {
	if w.closed {
		return
	}
	width, height = w.width, w.height
	return
}

func (w *Window) Show() {
	if w.closed {
		return
	}
	w.conn.MapWindow(w.id)
	if true {
		err := icccm.WmProtocolsSet(w.xu, w.id, []string{"WM_DELETE_WINDOW"})
		if err != nil {
			fmt.Println(err)
		}
	}
}

func (w *Window) Screen() (im draw.Image) {
	if w.closed {
		return
	}
	im = w.buffer
	return
}

func (w *Window) FlushImage() {
	if w.closed {
		return
	}
	xgraphics.PaintImg(w.xu, w.id, w.buffer)
}

func (w *Window) Close() (err error) {
	if w.closed {
		return
	}
	w.conn.DestroyWindow(w.id)
	w.closed = true
	return
}<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"strconv"

	"github.com/gorilla/websocket"
	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	cli "github.com/lxc/lxd/shared/cmd"
	"github.com/lxc/lxd/shared/i18n"
	"github.com/lxc/lxd/shared/logger"
	"github.com/lxc/lxd/shared/termios"
)

type cmdConsole struct {
	global *cmdGlobal

	flagShowLog bool
	flagType    string
}

func (c *cmdConsole) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = i18n.G("console [<remote>:]<instance>")
	cmd.Short = i18n.G("Attach to instance consoles")
	cmd.Long = cli.FormatSection(i18n.G("Description"), i18n.G(
		`Attach to instance consoles

This command allows you to interact with the boot console of an instance
as well as retrieve past log entries from it.`))

	cmd.RunE = c.Run
	cmd.Flags().BoolVar(&c.flagShowLog, "show-log", false, i18n.G("Retrieve the instance's console log"))
	cmd.Flags().StringVar(&c.flagType, "type", "console", i18n.G("Type of connection to establish: 'console' for serial console, 'vga' for SPICE graphical output"))

	return cmd
}

func (c *cmdConsole) sendTermSize(control *websocket.Conn) error {
	width, height, err := termios.GetSize(int(os.Stdout.Fd()))
	if err != nil {
		return err
	}

	logger.Debugf("Window size is now: %dx%d", width, height)

	msg := api.InstanceExecControl{}
	msg.Command = "window-resize"
	msg.Args = make(map[string]string)
	msg.Args["width"] = strconv.Itoa(width)
	msg.Args["height"] = strconv.Itoa(height)

	return control.WriteJSON(msg)
}

type readWriteCloser struct {
	io.Reader
	io.WriteCloser
}

type stdinMirror struct {
	r                 io.Reader
	consoleDisconnect chan<- bool
	foundEscape       *bool
}

// The pty has been switched to raw mode so we will only ever read a single
// byte. The buffer size is therefore uninteresting to us.
func (er stdinMirror) Read(p []byte) (int, error) {
	n, err := er.r.Read(p)

	v := rune(p[0])
	if v == '\u0001' && !*er.foundEscape {
		*er.foundEscape = true
		return 0, err
	}

	if v == 'q' && *er.foundEscape {
		select {
		case er.consoleDisconnect <- true:
			return 0, err
		default:
			return 0, err
		}
	}

	*er.foundEscape = false
	return n, err
}

func (c *cmdConsole) Run(cmd *cobra.Command, args []string) error {
	conf := c.global.conf

	// Sanity checks
	exit, err := c.global.CheckArgs(cmd, args, 1, 1)
	if exit {
		return err
	}

	// Validate flags.
	if !shared.StringInSlice(c.flagType, []string{"console", "vga"}) {
		return fmt.Errorf("Unknown output type %q", c.flagType)
	}

	// Connect to LXD
	remote, name, err := conf.ParseRemote(args[0])
	if err != nil {
		return err
	}

	d, err := conf.GetInstanceServer(remote)
	if err != nil {
		return err
	}

	// Show the current log if requested
	if c.flagShowLog {
		if c.flagType != "console" {
			return fmt.Errorf("The --show-log flag is only supported for by 'console' output type")
		}

		console := &lxd.InstanceConsoleLogArgs{}
		log, err := d.GetInstanceConsoleLog(name, console)
		if err != nil {
			return err
		}

		stuff, err := ioutil.ReadAll(log)
		if err != nil {
			return err
		}

		fmt.Printf("\n"+i18n.G("Console log:")+"\n\n%s\n", string(stuff))
		return nil
	}

	return c.Console(d, name)
}

func (c *cmdConsole) Console(d lxd.InstanceServer, name string) error {
	switch c.flagType {
	case "console":
		return c.console(d, name)
	case "vga":
		return c.vga(d, name)
	}
	return fmt.Errorf("Unknown console type %q", c.flagType)
}

func (c *cmdConsole) console(d lxd.InstanceServer, name string) error {
	// Configure the terminal
	cfd := int(os.Stdin.Fd())

	oldTTYstate, err := termios.MakeRaw(cfd)
	if err != nil {
		return err
	}
	defer termios.Restore(cfd, oldTTYstate)

	handler := c.controlSocketHandler

	var width, height int
	width, height, err = termios.GetSize(int(os.Stdin.Fd()))
	if err != nil {
		return err
	}

	// Prepare the remote console
	req := api.InstanceConsolePost{
		Width:  width,
		Height: height,
		Type:   "console",
	}

	consoleDisconnect := make(chan bool)
	sendDisconnect := make(chan bool)
	defer close(sendDisconnect)

	consoleArgs := lxd.InstanceConsoleArgs{
		Terminal: &readWriteCloser{stdinMirror{os.Stdin,
			sendDisconnect, new(bool)}, os.Stdout},
		Control:           handler,
		ConsoleDisconnect: consoleDisconnect,
	}

	go func() {
		<-sendDisconnect
		close(consoleDisconnect)
	}()

	fmt.Printf(i18n.G("To detach from the console, press: <ctrl>+a q") + "\n\r")

	// Attach to the instance console
	op, err := d.ConsoleInstance(name, req, &consoleArgs)
	if err != nil {
		return err
	}

	// Wait for the operation to complete
	err = op.Wait()
	if err != nil {
		return err
	}

	return nil
}

func (c *cmdConsole) vga(d lxd.InstanceServer, name string) error {
	conf := c.global.conf

	// We currently use the control websocket just to abort in case of errors.
	controlDone := make(chan struct{}, 1)
	handler := func(control *websocket.Conn) {
		<-controlDone
		closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
		control.WriteMessage(websocket.CloseMessage, closeMsg)
	}

	// Prepare the remote console.
	req := api.InstanceConsolePost{
		Type: "vga",
	}

	consoleDisconnect := make(chan bool)
	sendDisconnect := make(chan bool)
	defer close(sendDisconnect)

	consoleArgs := lxd.InstanceConsoleArgs{
		Control:           handler,
		ConsoleDisconnect: consoleDisconnect,
	}

	go func() {
		<-sendDisconnect
		close(consoleDisconnect)
	}()

	// Create a temporary unix socket mirroring the instance's spice socket.
	if !shared.PathExists(conf.ConfigPath("sockets")) {
		err := os.MkdirAll(conf.ConfigPath("sockets"), 0700)
		if err != nil {
			return err
		}
	}

	path, err := ioutil.TempFile(conf.ConfigPath("sockets"), "*.spice")
	if err != nil {
		return err
	}
	err = os.Remove(path.Name())
	if err != nil {
		return err
	}
	path.Close()

	socket := path.Name()
	listener, err := net.Listen("unix", socket)
	if err != nil {
		return err
	}
	defer listener.Close()
	defer os.Remove(path.Name())

	op, connect, err := d.ConsoleInstanceDynamic(name, req, &consoleArgs)
	if err != nil {
		return err
	}

	// Handle connections to the socket.
	go func() {
		count := 0

		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			count++

			go func(conn io.ReadWriteCloser) {
				err = connect(conn)
				if err != nil {
					sendDisconnect <- true
				}

				count--
				if count == 0 {
					sendDisconnect <- true
				}
			}(conn)
		}
	}()

	// Use either spicy or remote-viewer if available.
	spicy, err := exec.LookPath("spicy")
	if err == nil {
		_, err := shared.RunCommand(spicy, fmt.Sprintf("--uri=spice+unix://%s", socket))
		if err != nil {
			return err
		}
	} else {
		remoteViewer, err := exec.LookPath("remote-viewer")
		if err == nil {
			_, err := shared.RunCommand(remoteViewer, fmt.Sprintf("spice+unix://%s", socket))
			if err != nil {
				return err
			}
		} else {
			fmt.Println(i18n.G("LXD automatically uses either spicy or remote-viewer when present."))
			fmt.Println(i18n.G("As neither could be found, the raw SPICE socket can be found at:"))
			fmt.Printf("  %s\n", socket)
		}
	}

	// Wait for the operation to complete.
	err = op.Wait()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>lxc: Populate cmdConsole.flagType also when ran manually<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"strconv"

	"github.com/gorilla/websocket"
	"github.com/spf13/cobra"

	"github.com/lxc/lxd/client"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	cli "github.com/lxc/lxd/shared/cmd"
	"github.com/lxc/lxd/shared/i18n"
	"github.com/lxc/lxd/shared/logger"
	"github.com/lxc/lxd/shared/termios"
)

type cmdConsole struct {
	global *cmdGlobal

	flagShowLog bool
	flagType    string
}

func (c *cmdConsole) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = i18n.G("console [<remote>:]<instance>")
	cmd.Short = i18n.G("Attach to instance consoles")
	cmd.Long = cli.FormatSection(i18n.G("Description"), i18n.G(
		`Attach to instance consoles

This command allows you to interact with the boot console of an instance
as well as retrieve past log entries from it.`))

	cmd.RunE = c.Run
	cmd.Flags().BoolVar(&c.flagShowLog, "show-log", false, i18n.G("Retrieve the instance's console log"))
	cmd.Flags().StringVar(&c.flagType, "type", "console", i18n.G("Type of connection to establish: 'console' for serial console, 'vga' for SPICE graphical output"))

	return cmd
}

func (c *cmdConsole) sendTermSize(control *websocket.Conn) error {
	width, height, err := termios.GetSize(int(os.Stdout.Fd()))
	if err != nil {
		return err
	}

	logger.Debugf("Window size is now: %dx%d", width, height)

	msg := api.InstanceExecControl{}
	msg.Command = "window-resize"
	msg.Args = make(map[string]string)
	msg.Args["width"] = strconv.Itoa(width)
	msg.Args["height"] = strconv.Itoa(height)

	return control.WriteJSON(msg)
}

type readWriteCloser struct {
	io.Reader
	io.WriteCloser
}

type stdinMirror struct {
	r                 io.Reader
	consoleDisconnect chan<- bool
	foundEscape       *bool
}

// The pty has been switched to raw mode so we will only ever read a single
// byte. The buffer size is therefore uninteresting to us.
func (er stdinMirror) Read(p []byte) (int, error) {
	n, err := er.r.Read(p)

	v := rune(p[0])
	if v == '\u0001' && !*er.foundEscape {
		*er.foundEscape = true
		return 0, err
	}

	if v == 'q' && *er.foundEscape {
		select {
		case er.consoleDisconnect <- true:
			return 0, err
		default:
			return 0, err
		}
	}

	*er.foundEscape = false
	return n, err
}

func (c *cmdConsole) Run(cmd *cobra.Command, args []string) error {
	conf := c.global.conf

	// Sanity checks
	exit, err := c.global.CheckArgs(cmd, args, 1, 1)
	if exit {
		return err
	}

	// Validate flags.
	if !shared.StringInSlice(c.flagType, []string{"console", "vga"}) {
		return fmt.Errorf("Unknown output type %q", c.flagType)
	}

	// Connect to LXD
	remote, name, err := conf.ParseRemote(args[0])
	if err != nil {
		return err
	}

	d, err := conf.GetInstanceServer(remote)
	if err != nil {
		return err
	}

	// Show the current log if requested
	if c.flagShowLog {
		if c.flagType != "console" {
			return fmt.Errorf("The --show-log flag is only supported for by 'console' output type")
		}

		console := &lxd.InstanceConsoleLogArgs{}
		log, err := d.GetInstanceConsoleLog(name, console)
		if err != nil {
			return err
		}

		stuff, err := ioutil.ReadAll(log)
		if err != nil {
			return err
		}

		fmt.Printf("\n"+i18n.G("Console log:")+"\n\n%s\n", string(stuff))
		return nil
	}

	return c.Console(d, name)
}

func (c *cmdConsole) Console(d lxd.InstanceServer, name string) error {
	if c.flagType == "" {
		c.flagType = "console"
	}
	switch c.flagType {
	case "console":
		return c.console(d, name)
	case "vga":
		return c.vga(d, name)
	}
	return fmt.Errorf("Unknown console type %q", c.flagType)
}

func (c *cmdConsole) console(d lxd.InstanceServer, name string) error {
	// Configure the terminal
	cfd := int(os.Stdin.Fd())

	oldTTYstate, err := termios.MakeRaw(cfd)
	if err != nil {
		return err
	}
	defer termios.Restore(cfd, oldTTYstate)

	handler := c.controlSocketHandler

	var width, height int
	width, height, err = termios.GetSize(int(os.Stdin.Fd()))
	if err != nil {
		return err
	}

	// Prepare the remote console
	req := api.InstanceConsolePost{
		Width:  width,
		Height: height,
		Type:   "console",
	}

	consoleDisconnect := make(chan bool)
	sendDisconnect := make(chan bool)
	defer close(sendDisconnect)

	consoleArgs := lxd.InstanceConsoleArgs{
		Terminal: &readWriteCloser{stdinMirror{os.Stdin,
			sendDisconnect, new(bool)}, os.Stdout},
		Control:           handler,
		ConsoleDisconnect: consoleDisconnect,
	}

	go func() {
		<-sendDisconnect
		close(consoleDisconnect)
	}()

	fmt.Printf(i18n.G("To detach from the console, press: <ctrl>+a q") + "\n\r")

	// Attach to the instance console
	op, err := d.ConsoleInstance(name, req, &consoleArgs)
	if err != nil {
		return err
	}

	// Wait for the operation to complete
	err = op.Wait()
	if err != nil {
		return err
	}

	return nil
}

func (c *cmdConsole) vga(d lxd.InstanceServer, name string) error {
	conf := c.global.conf

	// We currently use the control websocket just to abort in case of errors.
	controlDone := make(chan struct{}, 1)
	handler := func(control *websocket.Conn) {
		<-controlDone
		closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
		control.WriteMessage(websocket.CloseMessage, closeMsg)
	}

	// Prepare the remote console.
	req := api.InstanceConsolePost{
		Type: "vga",
	}

	consoleDisconnect := make(chan bool)
	sendDisconnect := make(chan bool)
	defer close(sendDisconnect)

	consoleArgs := lxd.InstanceConsoleArgs{
		Control:           handler,
		ConsoleDisconnect: consoleDisconnect,
	}

	go func() {
		<-sendDisconnect
		close(consoleDisconnect)
	}()

	// Create a temporary unix socket mirroring the instance's spice socket.
	if !shared.PathExists(conf.ConfigPath("sockets")) {
		err := os.MkdirAll(conf.ConfigPath("sockets"), 0700)
		if err != nil {
			return err
		}
	}

	path, err := ioutil.TempFile(conf.ConfigPath("sockets"), "*.spice")
	if err != nil {
		return err
	}
	err = os.Remove(path.Name())
	if err != nil {
		return err
	}
	path.Close()

	socket := path.Name()
	listener, err := net.Listen("unix", socket)
	if err != nil {
		return err
	}
	defer listener.Close()
	defer os.Remove(path.Name())

	op, connect, err := d.ConsoleInstanceDynamic(name, req, &consoleArgs)
	if err != nil {
		return err
	}

	// Handle connections to the socket.
	go func() {
		count := 0

		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			count++

			go func(conn io.ReadWriteCloser) {
				err = connect(conn)
				if err != nil {
					sendDisconnect <- true
				}

				count--
				if count == 0 {
					sendDisconnect <- true
				}
			}(conn)
		}
	}()

	// Use either spicy or remote-viewer if available.
	spicy, err := exec.LookPath("spicy")
	if err == nil {
		_, err := shared.RunCommand(spicy, fmt.Sprintf("--uri=spice+unix://%s", socket))
		if err != nil {
			return err
		}
	} else {
		remoteViewer, err := exec.LookPath("remote-viewer")
		if err == nil {
			_, err := shared.RunCommand(remoteViewer, fmt.Sprintf("spice+unix://%s", socket))
			if err != nil {
				return err
			}
		} else {
			fmt.Println(i18n.G("LXD automatically uses either spicy or remote-viewer when present."))
			fmt.Println(i18n.G("As neither could be found, the raw SPICE socket can be found at:"))
			fmt.Printf("  %s\n", socket)
		}
	}

	// Wait for the operation to complete.
	err = op.Wait()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Package migration provides the primitives for migration in LXD.
//
// See https://github.com/lxc/lxd/blob/master/specs/migration.md for a complete
// description.

package main

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
	"google.golang.org/protobuf/proto"

	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/migration"
	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/idmap"
	"github.com/lxc/lxd/shared/logger"
)

type migrationFields struct {
	controlSecret string
	controlConn   *websocket.Conn
	controlLock   sync.Mutex

	criuSecret string
	criuConn   *websocket.Conn

	fsSecret string
	fsConn   *websocket.Conn

	// container specific fields
	live         bool
	instanceOnly bool
	instance     instance.Instance

	// storage specific fields
	volumeOnly bool
}

func (c *migrationFields) send(m proto.Message) error {
	/* gorilla websocket doesn't allow concurrent writes, and
	 * panic()s if it sees them (which is reasonable). If e.g. we
	 * happen to fail, get scheduled, start our write, then get
	 * unscheduled before the write is bit to a new thread which is
	 * receiving an error from the other side (due to our previous
	 * close), we can engage in these concurrent writes, which
	 * casuses the whole daemon to panic.
	 *
	 * Instead, let's lock sends to the controlConn so that we only ever
	 * write one message at the time.
	 */
	c.controlLock.Lock()
	defer c.controlLock.Unlock()

	err := migration.ProtoSend(c.controlConn, m)
	if err != nil {
		return err
	}

	return nil
}

func (c *migrationFields) recv(m proto.Message) error {
	return migration.ProtoRecv(c.controlConn, m)
}

func (c *migrationFields) disconnect() {
	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")

	c.controlLock.Lock()
	if c.controlConn != nil {
		c.controlConn.WriteMessage(websocket.CloseMessage, closeMsg)
		c.controlConn.Close()
		c.controlConn = nil /* don't close twice */
	}
	c.controlLock.Unlock()

	/* Below we just Close(), which doesn't actually write to the
	 * websocket, it just closes the underlying connection. If e.g. there
	 * is still a filesystem transfer going on, but the other side has run
	 * out of disk space, writing an actual CloseMessage here will cause
	 * gorilla websocket to panic. Instead, we just force close this
	 * connection, since we report the error over the control channel
	 * anyway.
	 */
	if c.fsConn != nil {
		c.fsConn.Close()
	}

	if c.criuConn != nil {
		c.criuConn.Close()
	}
}

func (c *migrationFields) sendControl(err error) {
	c.controlLock.Lock()
	if c.controlConn != nil {
		migration.ProtoSendControl(c.controlConn, err)
	}
	c.controlLock.Unlock()

	if err != nil {
		c.disconnect()
	}
}

func (c *migrationFields) controlChannel() <-chan *migration.MigrationControl {
	ch := make(chan *migration.MigrationControl)
	go func() {
		msg := migration.MigrationControl{}
		err := c.recv(&msg)
		if err != nil {
			logger.Debugf("Got error reading migration control socket %s", err)
			close(ch)
			return
		}
		ch <- &msg
	}()

	return ch
}

type migrationSourceWs struct {
	migrationFields

	allConnected chan struct{}
}

func (s *migrationSourceWs) Metadata() interface{} {
	secrets := shared.Jmap{
		"control": s.controlSecret,
		"fs":      s.fsSecret,
	}

	if s.criuSecret != "" {
		secrets["criu"] = s.criuSecret
	}

	return secrets
}

func (s *migrationSourceWs) Connect(op *operations.Operation, r *http.Request, w http.ResponseWriter) error {
	secret := r.FormValue("secret")
	if secret == "" {
		return fmt.Errorf("missing secret")
	}

	var conn **websocket.Conn

	switch secret {
	case s.controlSecret:
		conn = &s.controlConn
	case s.criuSecret:
		conn = &s.criuConn
	case s.fsSecret:
		conn = &s.fsConn
	default:
		// If we didn't find the right secret, the user provided a bad
		// one, which 403, not 404, since this operation actually
		// exists.
		return os.ErrPermission
	}

	c, err := shared.WebsocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		return err
	}

	*conn = c

	// Check criteria for considering all channels to be connected.
	if s.instance != nil && s.instance.Type() == instancetype.Container && s.live && s.criuConn == nil {
		return nil
	}

	if s.controlConn == nil {
		return nil
	}

	if s.fsConn == nil {
		return nil
	}

	close(s.allConnected)

	return nil
}

func (s *migrationSourceWs) ConnectTarget(certificate string, operation string, websockets map[string]string) error {
	var err error
	var cert *x509.Certificate

	if certificate != "" {
		certBlock, _ := pem.Decode([]byte(certificate))
		if certBlock == nil {
			return fmt.Errorf("Invalid certificate")
		}

		cert, err = x509.ParseCertificate(certBlock.Bytes)
		if err != nil {
			return err
		}
	}

	config, err := shared.GetTLSConfig("", "", "", cert)
	if err != nil {
		return err
	}

	dialer := websocket.Dialer{
		TLSClientConfig: config,
		NetDial:         shared.RFC3493Dialer,
	}

	for name, secret := range websockets {
		var conn **websocket.Conn

		switch name {
		case "control":
			conn = &s.controlConn
		case "fs":
			conn = &s.fsConn
		case "criu":
			conn = &s.criuConn
		default:
			return fmt.Errorf("Unknown secret provided: %s", name)
		}

		query := url.Values{"secret": []string{secret}}

		// The URL is a https URL to the operation, mangle to be a wss URL to the secret
		wsUrl := fmt.Sprintf("wss://%s/websocket?%s", strings.TrimPrefix(operation, "https://"), query.Encode())

		wsConn, _, err := dialer.Dial(wsUrl, http.Header{})
		if err != nil {
			return err
		}

		*conn = wsConn
	}

	close(s.allConnected)

	return nil
}

type migrationSink struct {
	// We are pulling the entity from src in pull mode.
	src migrationFields
	// The entity is pushed from src to dest in push mode. Note that
	// websocket connections are not set in push mode. Only the secret
	// fields are used since the client will connect to the sockets.
	dest migrationFields

	url          string
	dialer       websocket.Dialer
	allConnected chan struct{}
	push         bool
	refresh      bool
}

type MigrationSinkArgs struct {
	// General migration fields
	Dialer  websocket.Dialer
	Push    bool
	Secrets map[string]string
	Url     string

	// Instance specific fields
	Instance     instance.Instance
	InstanceOnly bool
	Idmap        *idmap.IdmapSet
	Live         bool
	Refresh      bool
	Snapshots    []*migration.Snapshot

	// Storage specific fields
	VolumeOnly bool
	VolumeSize int64

	// Transport specific fields
	RsyncFeatures []string
}

func (c *migrationSink) connectWithSecret(secret string) (*websocket.Conn, error) {
	query := url.Values{"secret": []string{secret}}

	// The URL is a https URL to the operation, mangle to be a wss URL to the secret
	wsUrl := fmt.Sprintf("wss://%s/websocket?%s", strings.TrimPrefix(c.url, "https://"), query.Encode())

	conn, _, err := c.dialer.Dial(wsUrl, http.Header{})
	if err != nil {
		return nil, err
	}

	return conn, err
}

func (s *migrationSink) Metadata() interface{} {
	secrets := shared.Jmap{
		"control": s.dest.controlSecret,
		"fs":      s.dest.fsSecret,
	}

	if s.dest.criuSecret != "" {
		secrets["criu"] = s.dest.criuSecret
	}

	return secrets
}

func (s *migrationSink) Connect(op *operations.Operation, r *http.Request, w http.ResponseWriter) error {
	secret := r.FormValue("secret")
	if secret == "" {
		return fmt.Errorf("missing secret")
	}

	var conn **websocket.Conn

	switch secret {
	case s.dest.controlSecret:
		conn = &s.dest.controlConn
	case s.dest.criuSecret:
		conn = &s.dest.criuConn
	case s.dest.fsSecret:
		conn = &s.dest.fsConn
	default:
		/* If we didn't find the right secret, the user provided a bad one,
		 * which 403, not 404, since this operation actually exists */
		return os.ErrPermission
	}

	c, err := shared.WebsocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		return err
	}

	*conn = c

	// Check criteria for considering all channels to be connected.
	if s.src.instance != nil && s.src.instance.Type() == instancetype.Container && s.dest.live && s.dest.criuConn == nil {
		return nil
	}

	if s.dest.controlConn == nil {
		return nil
	}

	if s.dest.fsConn == nil {
		return nil
	}

	close(s.allConnected)

	return nil
}
<commit_msg>lxd/migrate: Updates migrationSourceWs to use util.ExtractTCPConn and util.SetTCPTimeouts<commit_after>// Package migration provides the primitives for migration in LXD.
//
// See https://github.com/lxc/lxd/blob/master/specs/migration.md for a complete
// description.

package main

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
	"google.golang.org/protobuf/proto"

	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/migration"
	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/idmap"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
)

type migrationFields struct {
	controlSecret string
	controlConn   *websocket.Conn
	controlLock   sync.Mutex

	criuSecret string
	criuConn   *websocket.Conn

	fsSecret string
	fsConn   *websocket.Conn

	// container specific fields
	live         bool
	instanceOnly bool
	instance     instance.Instance

	// storage specific fields
	volumeOnly bool
}

func (c *migrationFields) send(m proto.Message) error {
	/* gorilla websocket doesn't allow concurrent writes, and
	 * panic()s if it sees them (which is reasonable). If e.g. we
	 * happen to fail, get scheduled, start our write, then get
	 * unscheduled before the write is bit to a new thread which is
	 * receiving an error from the other side (due to our previous
	 * close), we can engage in these concurrent writes, which
	 * casuses the whole daemon to panic.
	 *
	 * Instead, let's lock sends to the controlConn so that we only ever
	 * write one message at the time.
	 */
	c.controlLock.Lock()
	defer c.controlLock.Unlock()

	err := migration.ProtoSend(c.controlConn, m)
	if err != nil {
		return err
	}

	return nil
}

func (c *migrationFields) recv(m proto.Message) error {
	return migration.ProtoRecv(c.controlConn, m)
}

func (c *migrationFields) disconnect() {
	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")

	c.controlLock.Lock()
	if c.controlConn != nil {
		c.controlConn.WriteMessage(websocket.CloseMessage, closeMsg)
		c.controlConn.Close()
		c.controlConn = nil /* don't close twice */
	}
	c.controlLock.Unlock()

	/* Below we just Close(), which doesn't actually write to the
	 * websocket, it just closes the underlying connection. If e.g. there
	 * is still a filesystem transfer going on, but the other side has run
	 * out of disk space, writing an actual CloseMessage here will cause
	 * gorilla websocket to panic. Instead, we just force close this
	 * connection, since we report the error over the control channel
	 * anyway.
	 */
	if c.fsConn != nil {
		c.fsConn.Close()
	}

	if c.criuConn != nil {
		c.criuConn.Close()
	}
}

func (c *migrationFields) sendControl(err error) {
	c.controlLock.Lock()
	if c.controlConn != nil {
		migration.ProtoSendControl(c.controlConn, err)
	}
	c.controlLock.Unlock()

	if err != nil {
		c.disconnect()
	}
}

func (c *migrationFields) controlChannel() <-chan *migration.MigrationControl {
	ch := make(chan *migration.MigrationControl)
	go func() {
		msg := migration.MigrationControl{}
		err := c.recv(&msg)
		if err != nil {
			logger.Debugf("Got error reading migration control socket %s", err)
			close(ch)
			return
		}
		ch <- &msg
	}()

	return ch
}

type migrationSourceWs struct {
	migrationFields

	allConnected chan struct{}
}

func (s *migrationSourceWs) Metadata() interface{} {
	secrets := shared.Jmap{
		"control": s.controlSecret,
		"fs":      s.fsSecret,
	}

	if s.criuSecret != "" {
		secrets["criu"] = s.criuSecret
	}

	return secrets
}

func (s *migrationSourceWs) Connect(op *operations.Operation, r *http.Request, w http.ResponseWriter) error {
	secret := r.FormValue("secret")
	if secret == "" {
		return fmt.Errorf("missing secret")
	}

	var conn **websocket.Conn

	switch secret {
	case s.controlSecret:
		conn = &s.controlConn
	case s.criuSecret:
		conn = &s.criuConn
	case s.fsSecret:
		conn = &s.fsConn
	default:
		// If we didn't find the right secret, the user provided a bad
		// one, which 403, not 404, since this operation actually
		// exists.
		return os.ErrPermission
	}

	c, err := shared.WebsocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		return err
	}

	remoteTCP, err := util.ExtractTCPConn(c.UnderlyingConn())
	if err != nil {
		logger.Error("Failed extracting TCP connection from remote connection", log.Ctx{"err": err})
	} else {
		err := util.SetTCPTimeouts(remoteTCP)
		if err != nil {
			logger.Error("Failed setting TCP timeouts on remote connection", log.Ctx{"err": err})
		}
	}

	*conn = c

	// Check criteria for considering all channels to be connected.
	if s.instance != nil && s.instance.Type() == instancetype.Container && s.live && s.criuConn == nil {
		return nil
	}

	if s.controlConn == nil {
		return nil
	}

	if s.fsConn == nil {
		return nil
	}

	close(s.allConnected)

	return nil
}

func (s *migrationSourceWs) ConnectTarget(certificate string, operation string, websockets map[string]string) error {
	var err error
	var cert *x509.Certificate

	if certificate != "" {
		certBlock, _ := pem.Decode([]byte(certificate))
		if certBlock == nil {
			return fmt.Errorf("Invalid certificate")
		}

		cert, err = x509.ParseCertificate(certBlock.Bytes)
		if err != nil {
			return err
		}
	}

	config, err := shared.GetTLSConfig("", "", "", cert)
	if err != nil {
		return err
	}

	dialer := websocket.Dialer{
		TLSClientConfig: config,
		NetDial:         shared.RFC3493Dialer,
	}

	for name, secret := range websockets {
		var conn **websocket.Conn

		switch name {
		case "control":
			conn = &s.controlConn
		case "fs":
			conn = &s.fsConn
		case "criu":
			conn = &s.criuConn
		default:
			return fmt.Errorf("Unknown secret provided: %s", name)
		}

		query := url.Values{"secret": []string{secret}}

		// The URL is a https URL to the operation, mangle to be a wss URL to the secret
		wsUrl := fmt.Sprintf("wss://%s/websocket?%s", strings.TrimPrefix(operation, "https://"), query.Encode())

		wsConn, _, err := dialer.Dial(wsUrl, http.Header{})
		if err != nil {
			return err
		}

		*conn = wsConn
	}

	close(s.allConnected)

	return nil
}

type migrationSink struct {
	// We are pulling the entity from src in pull mode.
	src migrationFields
	// The entity is pushed from src to dest in push mode. Note that
	// websocket connections are not set in push mode. Only the secret
	// fields are used since the client will connect to the sockets.
	dest migrationFields

	url          string
	dialer       websocket.Dialer
	allConnected chan struct{}
	push         bool
	refresh      bool
}

type MigrationSinkArgs struct {
	// General migration fields
	Dialer  websocket.Dialer
	Push    bool
	Secrets map[string]string
	Url     string

	// Instance specific fields
	Instance     instance.Instance
	InstanceOnly bool
	Idmap        *idmap.IdmapSet
	Live         bool
	Refresh      bool
	Snapshots    []*migration.Snapshot

	// Storage specific fields
	VolumeOnly bool
	VolumeSize int64

	// Transport specific fields
	RsyncFeatures []string
}

func (c *migrationSink) connectWithSecret(secret string) (*websocket.Conn, error) {
	query := url.Values{"secret": []string{secret}}

	// The URL is a https URL to the operation, mangle to be a wss URL to the secret
	wsUrl := fmt.Sprintf("wss://%s/websocket?%s", strings.TrimPrefix(c.url, "https://"), query.Encode())

	conn, _, err := c.dialer.Dial(wsUrl, http.Header{})
	if err != nil {
		return nil, err
	}

	return conn, err
}

func (s *migrationSink) Metadata() interface{} {
	secrets := shared.Jmap{
		"control": s.dest.controlSecret,
		"fs":      s.dest.fsSecret,
	}

	if s.dest.criuSecret != "" {
		secrets["criu"] = s.dest.criuSecret
	}

	return secrets
}

func (s *migrationSink) Connect(op *operations.Operation, r *http.Request, w http.ResponseWriter) error {
	secret := r.FormValue("secret")
	if secret == "" {
		return fmt.Errorf("missing secret")
	}

	var conn **websocket.Conn

	switch secret {
	case s.dest.controlSecret:
		conn = &s.dest.controlConn
	case s.dest.criuSecret:
		conn = &s.dest.criuConn
	case s.dest.fsSecret:
		conn = &s.dest.fsConn
	default:
		/* If we didn't find the right secret, the user provided a bad one,
		 * which 403, not 404, since this operation actually exists */
		return os.ErrPermission
	}

	c, err := shared.WebsocketUpgrader.Upgrade(w, r, nil)
	if err != nil {
		return err
	}

	*conn = c

	// Check criteria for considering all channels to be connected.
	if s.src.instance != nil && s.src.instance.Type() == instancetype.Container && s.dest.live && s.dest.criuConn == nil {
		return nil
	}

	if s.dest.controlConn == nil {
		return nil
	}

	if s.dest.fsConn == nil {
		return nil
	}

	close(s.allConnected)

	return nil
}
<|endoftext|>