<commit_before>package api

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
)

const maxBodyLen = 2 << 20 // 2 MiB

// contextKey is a value for use with context.WithValue.
type contextKey struct {
	name string
}

// UserContextKey is a context key. It can be used in
// HTTP handlers to attach a client session to a request.
//
// The session is then read by the Transport and used
// for authentication.
//
// The associated value will be of *User type.
var UserContextKey = &contextKey{"client-session"}

// UserFunc is used to extract a user, which requires authentication,
// from a request.
type UserFunc func(*http.Request) *User

// User represents a user, which is going to be authenticated.
type User struct {
	Username string // name of the user
	Team     string // team which the user belongs to
}

// WithRequest attaches a user to the given request.
func (u *User) WithRequest(req *http.Request) *http.Request {
	return req.WithContext(context.WithValue(req.Context(), UserContextKey, u))
}

// Strings implements the fmt.Stringer interface.
func (u *User) String() string {
	return u.Team + "/" + u.Username
}

// Session represents a user session.
type Session struct {
	ClientID string // an ID of an authenticated session.
	User     *User  // authenticated user
}

// Valid implements the stack.Validator interface.
func (s *Session) Valid() error {
	if s == nil {
		return errors.New("session is missing")
	}
	if s.ClientID == "" {
		return errors.New("empty client ID")
	}
	return nil
}

func (s *Session) writeTo(req *http.Request) {
	req.Header.Set("Authorization", "Bearer "+s.ClientID)
}

func (s *Session) readFrom(req *http.Request) {
	if auth := req.Header.Get("Authorization"); strings.HasPrefix(auth, "Bearer ") {
		s.ClientID = auth[len("Bearer "):]
	}
}

// AuthFunc is used to fetch session information.
//
// When Refresh is true, AuthFunc must obtain
// a session directly from an authorisation
// endpoint, in case it was cached.
type AuthFunc func(*AuthOptions) (*Session, error)

// AuthOptions represents the arguments for AuthFunc.
type AuthOptions struct {
	// User represents a user, which we are building
	// authorisation for.
	User *User

	// Refresh requests to refresh a session when true.
	Refresh bool

	// Request is a copy of the request, which
	// we're building authorisation for.
	Request *http.Request

	// Error is the error that was the reason
	// of temporary failure received while trying
	// to sending request.
	//
	// Error is non-nil when request failed with
	// temporary error and Transport is retrying.
	Error error
}

type httpTransport interface {
	http.RoundTripper
	httpRequestCanceler
	httpIdleConnectionsCloser
}

type httpRequestCanceler interface {
	CancelRequest(*http.Request)
}

type httpIdleConnectionsCloser interface {
	CloseIdleConnections()
}

// nopCloser is like ioutil.NopCloser, but wraps io.ReadSeeker.
type nopCloser struct{ io.ReadSeeker }

func (nopCloser) Close() (_ error) { return }

var _ io.ReadSeeker = nopCloser{bytes.NewReader(nil)}

// Transport is a signing transport that
// authorizes each request with jSession.clientId.
type Transport struct {
	http.RoundTripper          // a transport to use; optional, by default http.DefaultClient.Transport
	AuthFunc          AuthFunc // clientID authorisation to use; required
	UserFunc          UserFunc // used to get a user from request to authenticate for; by default looked up with UserContextKey
	Host              string   // original Host name, overwrites req.Host; optional
	RetryNum          int      // number of retries in case of temporary failures; optional, by default 3

	// request mapping implementation stolen from:
	//
	//   https://github.com/golang/oauth2/blob/master/transport.go
	//
	mu     sync.Mutex                      // guards modReq
	modReq map[*http.Request]*http.Request // original -> modified
}

var _ httpTransport = (*Transport)(nil)

// NewSingleUser returns a HTTP transport that authenticates all
// requests for the given user.
//
// This is usefull for api clients that give no control over HTTP requests,
// like the swagger-generated ones (remote.api).
func (t *Transport) NewSingleUser(u *User) http.RoundTripper {
	return &Transport{
		RoundTripper: t.RoundTripper,
		AuthFunc:     t.AuthFunc,
		UserFunc:     func(*http.Request) *User { return u },
		Host:         t.Host,
		RetryNum:     t.RetryNum,
	}
}

// RoundTrip implements the http.RoundTripper interface.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	user := t.user(req)
	if user == nil {
		return t.roundTripper().RoundTrip(req)
	}

	reqCopy := copyRequest(req) // per RoundTripper contract
	t.setModReq(req, reqCopy)

	switch reqCopy.Body.(type) {
	case nil:
	case io.ReadSeeker:
	default:
		p, err := ioutil.ReadAll(io.LimitReader(reqCopy.Body, maxBodyLen))
		if err != nil {
			return nil, err // non-retryable
		}

		reqCopy.Body = nopCloser{bytes.NewReader(p)}
	}

	var refresh bool
	var lastErr error
	var session *Session

	for attempts := t.retryNum(); attempts >= 0; attempts-- {
		if reqCopy.Body != nil {
			if _, err := reqCopy.Body.(io.ReadSeeker).Seek(0, io.SeekStart); err != nil {
				return nil, err // non-retryable
			}
		}

		opts := &AuthOptions{
			User:    user,
			Refresh: refresh,
			Request: reqCopy,
			Error:   lastErr,
		}

		s, err := t.AuthFunc(opts)
		if err != nil {
			return nil, err // non-retryable
		}

		if err := s.Valid(); err != nil {
			return nil, err // non-retryable
		}

		session = s
		refresh = false

		session.writeTo(reqCopy)

		if t.Host != "" {
			reqCopy.Host = t.Host
		}

		resp, err := t.roundTripper().RoundTrip(reqCopy)
		if e, ok := err.(net.Error); ok && (e.Temporary() || e.Timeout()) {
			lastErr = err
			continue // retry
		}

		if err != nil {
			t.setModReq(req, nil)
			return nil, err // non-retryable
		}

		if resp.StatusCode == http.StatusUnauthorized {
			resp.Body.Close()
			lastErr = errors.New(http.StatusText(resp.StatusCode))
			refresh = true
			continue // retry and force new session
		}

		resp.Body = &onEOFReader{
			rc: resp.Body,
			fn: func() { t.setModReq(req, nil) },
		}

		return resp, nil
	}

	t.setModReq(req, nil)

	return nil, fmt.Errorf("api: error sending request: %v", lastErr)
}

func (t *Transport) CancelRequest(req *http.Request) {
	if rc, ok := t.roundTripper().(httpRequestCanceler); ok {
		t.mu.Lock()
		modReq := t.modReq[req]
		delete(t.modReq, req)
		t.mu.Unlock()

		rc.CancelRequest(modReq)
	}
}

func (t *Transport) CloseIdleConnections() {
	if icl, ok := t.roundTripper().(httpIdleConnectionsCloser); ok {
		icl.CloseIdleConnections()
	}
}

func (t *Transport) user(req *http.Request) *User {
	if t.UserFunc != nil {
		return t.UserFunc(req)
	}
	user, _ := req.Context().Value(UserContextKey).(*User)
	return user
}

func (t *Transport) roundTripper() http.RoundTripper {
	if t.RoundTripper != nil {
		return t.RoundTripper
	}

	if http.DefaultClient.Transport != nil {
		return http.DefaultClient.Transport
	}

	return http.DefaultTransport
}

func (t *Transport) retryNum() int {
	if t.RetryNum > 0 {
		return t.RetryNum
	}
	return 3
}

func (t *Transport) setModReq(orig, mod *http.Request) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.modReq == nil {
		t.modReq = make(map[*http.Request]*http.Request)
	}
	if mod == nil {
		delete(t.modReq, orig)
	} else {
		t.modReq[orig] = mod
	}
}

func copyURL(u *url.URL) *url.URL {
	uCopy := *u
	if u.User != nil {
		userCopy := *u.User
		uCopy.User = &userCopy
	}

	return &uCopy
}

func copyRequest(req *http.Request) *http.Request {
	reqCopy := new(http.Request)
	*reqCopy = *req
	reqCopy.URL = copyURL(req.URL)
	reqCopy.Header = make(http.Header, len(req.Header))
	for k, s := range req.Header {
		reqCopy.Header[k] = append([]string(nil), s...)
	}
	return reqCopy
}

type onEOFReader struct {
	rc io.ReadCloser
	fn func()
}

func (r *onEOFReader) Read(p []byte) (n int, err error) {
	n, err = r.rc.Read(p)
	if err == io.EOF {
		r.runFunc()
	}
	return
}

func (r *onEOFReader) Close() error {
	err := r.rc.Close()
	r.runFunc()
	return err
}

func (r *onEOFReader) runFunc() {
	if fn := r.fn; fn != nil {
		fn()
		r.fn = nil
	}
}
<commit_msg>api: add debug logging<commit_after>package api

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/koding/logging"
)

const maxBodyLen = 2 << 20 // 2 MiB

var defaultLog = logging.NewCustom("api-transport", false)

// contextKey is a value for use with context.WithValue.
type contextKey struct {
	name string
}

// UserContextKey is a context key. It can be used in
// HTTP handlers to attach a client session to a request.
//
// The session is then read by the Transport and used
// for authentication.
//
// The associated value will be of *User type.
var UserContextKey = &contextKey{"client-session"}

// UserFunc is used to extract a user, which requires authentication,
// from a request.
type UserFunc func(*http.Request) *User

// User represents a user, which is going to be authenticated.
type User struct {
	Username string // name of the user
	Team     string // team which the user belongs to
}

// WithRequest attaches a user to the given request.
func (u *User) WithRequest(req *http.Request) *http.Request {
	return req.WithContext(context.WithValue(req.Context(), UserContextKey, u))
}

// Strings implements the fmt.Stringer interface.
func (u *User) String() string {
	return u.Team + "/" + u.Username
}

// Session represents a user session.
type Session struct {
	ClientID string // an ID of an authenticated session.
	User     *User  // authenticated user
}

// Valid implements the stack.Validator interface.
func (s *Session) Valid() error {
	if s == nil {
		return errors.New("session is missing")
	}
	if s.ClientID == "" {
		return errors.New("empty client ID")
	}
	return nil
}

func (s *Session) writeTo(req *http.Request) {
	req.Header.Set("Authorization", "Bearer "+s.ClientID)
}

func (s *Session) readFrom(req *http.Request) {
	if auth := req.Header.Get("Authorization"); strings.HasPrefix(auth, "Bearer ") {
		s.ClientID = auth[len("Bearer "):]
	}
}

// AuthFunc is used to fetch session information.
//
// When Refresh is true, AuthFunc must obtain
// a session directly from an authorisation
// endpoint, in case it was cached.
type AuthFunc func(*AuthOptions) (*Session, error)

// AuthOptions represents the arguments for AuthFunc.
type AuthOptions struct {
	// User represents a user, which we are building
	// authorisation for.
	User *User

	// Refresh requests to refresh a session when true.
	Refresh bool

	// Request is a copy of the request, which
	// we're building authorisation for.
	Request *http.Request

	// Error is the error that was the reason
	// of temporary failure received while trying
	// to sending request.
	//
	// Error is non-nil when request failed with
	// temporary error and Transport is retrying.
	Error error
}

type httpTransport interface {
	http.RoundTripper
	httpRequestCanceler
	httpIdleConnectionsCloser
}

type httpRequestCanceler interface {
	CancelRequest(*http.Request)
}

type httpIdleConnectionsCloser interface {
	CloseIdleConnections()
}

// nopCloser is like ioutil.NopCloser, but wraps io.ReadSeeker.
type nopCloser struct{ io.ReadSeeker }

func (nopCloser) Close() (_ error) { return }

var _ io.ReadSeeker = nopCloser{bytes.NewReader(nil)}

// Transport is a signing transport that
// authorizes each request with jSession.clientId.
type Transport struct {
	http.RoundTripper                // a transport to use; optional, by default http.DefaultClient.Transport
	AuthFunc          AuthFunc       // clientID authorisation to use; required
	UserFunc          UserFunc       // used to get a user from request to authenticate for; by default looked up with UserContextKey
	Host              string         // original Host name, overwrites req.Host; optional
	RetryNum          int            // number of retries in case of temporary failures; optional, by default 3
	Log               logging.Logger // logger to use

	// request mapping implementation stolen from:
	//
	//   https://github.com/golang/oauth2/blob/master/transport.go
	//
	mu     sync.Mutex                      // guards modReq
	modReq map[*http.Request]*http.Request // original -> modified
}

var _ httpTransport = (*Transport)(nil)

// NewSingleUser returns a HTTP transport that authenticates all
// requests for the given user.
//
// This is usefull for api clients that give no control over HTTP requests,
// like the swagger-generated ones (remote.api).
func (t *Transport) NewSingleUser(u *User) http.RoundTripper {
	return &Transport{
		RoundTripper: t.RoundTripper,
		AuthFunc:     t.AuthFunc,
		UserFunc:     func(*http.Request) *User { return u },
		Host:         t.Host,
		RetryNum:     t.RetryNum,
	}
}

// RoundTrip implements the http.RoundTripper interface.
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
	user := t.user(req)
	if user == nil {
		return t.roundTripper().RoundTrip(req)
	}

	reqCopy := copyRequest(req) // per RoundTripper contract
	t.setModReq(req, reqCopy)

	switch reqCopy.Body.(type) {
	case nil:
	case io.ReadSeeker:
	default:
		p, err := ioutil.ReadAll(io.LimitReader(reqCopy.Body, maxBodyLen))
		if err != nil {
			return nil, err // non-retryable
		}

		reqCopy.Body = nopCloser{bytes.NewReader(p)}
	}

	var refresh bool
	var lastErr error
	var session *Session

	for attempts := t.retryNum(); attempts >= 0; attempts-- {
		if reqCopy.Body != nil {
			if _, err := reqCopy.Body.(io.ReadSeeker).Seek(0, io.SeekStart); err != nil {
				return nil, err // non-retryable
			}
		}

		opts := &AuthOptions{
			User:    user,
			Refresh: refresh,
			Request: reqCopy,
			Error:   lastErr,
		}

		s, err := t.AuthFunc(opts)

		t.log().Debug("AuthFunc(%+v): session=%+v, error=%v", opts, s, err)

		if err != nil {
			return nil, err // non-retryable
		}

		if err := s.Valid(); err != nil {
			return nil, err // non-retryable
		}

		session = s
		refresh = false

		session.writeTo(reqCopy)

		if t.Host != "" {
			reqCopy.Host = t.Host
		}

		resp, err := t.roundTripper().RoundTrip(reqCopy)
		if e, ok := err.(net.Error); ok && (e.Temporary() || e.Timeout()) {
			lastErr = err
			continue // retry
		}

		if err != nil {
			t.setModReq(req, nil)
			return nil, err // non-retryable
		}

		if resp.StatusCode == http.StatusUnauthorized {
			resp.Body.Close()
			lastErr = errors.New(http.StatusText(resp.StatusCode))
			refresh = true
			continue // retry and force new session
		}

		resp.Body = &onEOFReader{
			rc: resp.Body,
			fn: func() { t.setModReq(req, nil) },
		}

		return resp, nil
	}

	t.setModReq(req, nil)

	return nil, fmt.Errorf("api: error sending request: %v", lastErr)
}

func (t *Transport) CancelRequest(req *http.Request) {
	if rc, ok := t.roundTripper().(httpRequestCanceler); ok {
		t.mu.Lock()
		modReq := t.modReq[req]
		delete(t.modReq, req)
		t.mu.Unlock()

		rc.CancelRequest(modReq)
	}
}

func (t *Transport) CloseIdleConnections() {
	if icl, ok := t.roundTripper().(httpIdleConnectionsCloser); ok {
		icl.CloseIdleConnections()
	}
}

func (t *Transport) user(req *http.Request) *User {
	if t.UserFunc != nil {
		return t.UserFunc(req)
	}
	user, _ := req.Context().Value(UserContextKey).(*User)
	return user
}

func (t *Transport) roundTripper() http.RoundTripper {
	if t.RoundTripper != nil {
		return t.RoundTripper
	}

	if http.DefaultClient.Transport != nil {
		return http.DefaultClient.Transport
	}

	return http.DefaultTransport
}

func (t *Transport) retryNum() int {
	if t.RetryNum > 0 {
		return t.RetryNum
	}
	return 3
}

func (t *Transport) log() logging.Logger {
	if t.Log != nil {
		return t.Log
	}
	return defaultLog
}

func (t *Transport) setModReq(orig, mod *http.Request) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.modReq == nil {
		t.modReq = make(map[*http.Request]*http.Request)
	}
	if mod == nil {
		delete(t.modReq, orig)
	} else {
		t.modReq[orig] = mod
	}
}

func copyURL(u *url.URL) *url.URL {
	uCopy := *u
	if u.User != nil {
		userCopy := *u.User
		uCopy.User = &userCopy
	}

	return &uCopy
}

func copyRequest(req *http.Request) *http.Request {
	reqCopy := new(http.Request)
	*reqCopy = *req
	reqCopy.URL = copyURL(req.URL)
	reqCopy.Header = make(http.Header, len(req.Header))
	for k, s := range req.Header {
		reqCopy.Header[k] = append([]string(nil), s...)
	}
	return reqCopy
}

type onEOFReader struct {
	rc io.ReadCloser
	fn func()
}

func (r *onEOFReader) Read(p []byte) (n int, err error) {
	n, err = r.rc.Read(p)
	if err == io.EOF {
		r.runFunc()
	}
	return
}

func (r *onEOFReader) Close() error {
	err := r.rc.Close()
	r.runFunc()
	return err
}

func (r *onEOFReader) runFunc() {
	if fn := r.fn; fn != nil {
		fn()
		r.fn = nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"koding/broker/cache"
	"koding/tools/config"
	"koding/tools/sockjs"
	"strconv"
	"strings"
	"time"

	"github.com/streadway/amqp"
)

type Client struct {
	Session        *sockjs.Session
	ControlChannel *amqp.Channel
	SocketId       string
	Broker         *Broker
	LastPayload    string
	Subscriptions  *cache.SubscriptionStorage
}

// NewClient retuns a new client that is defined on a given session.
func NewClient(session *sockjs.Session, broker *Broker) *Client {
	socketID := randomString()
	session.Tag = socketID

	controlChannel, err := broker.PublishConn.Channel()
	if err != nil {
		panic(err)
	}

	subscriptions, err := cache.NewStorage(STORAGE_BACKEND, socketID)
	if err != nil {
		panic(err)
	}

	return &Client{
		Session:        session,
		SocketId:       socketID,
		ControlChannel: controlChannel,
		Broker:         broker,
		Subscriptions:  subscriptions,
	}
}

// Close should be called whenever a client disconnects.
func (c *Client) Close() {
	log.Debug("Client Close Request for socketID: %v", c.SocketId)
	c.Subscriptions.Each(func(routingKeyPrefix interface{}) bool {
		c.RemoveFromRoute(routingKeyPrefix.(string))
		return true
	})

	c.Subscriptions.ClearWithTimeout()

	for {
		err := c.ControlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(c.SocketId)})
		if err == nil {
			break
		}
		if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != amqp.ChannelError {
			panic(err)
		}
		c.resetControlChannel()
	}

	log.Debug("Closing control channel for socketID: %v", c.SocketId)
	c.ControlChannel.Close()
}

// handleSessionMessage handles the received message from the client. It
// passes a response back to the client or publish the received message to a
// rabbitmq exchange for further process.
func (c *Client) handleSessionMessage(data interface{}) {
	defer log.RecoverAndLog()

	message := data.(map[string]interface{})
	log.Debug("Received message: %v", message)

	action := message["action"]
	switch action {
	case "subscribe":
		for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
			if err := c.Subscribe(routingKeyPrefix); err != nil {
				log.Error(err.Error())
			}
		}

		sendToClient(c.Session, "broker.subscribed", message["routingKeyPrefix"])

	case "resubscribe":
		globalMapMutex.Lock()
		defer globalMapMutex.Unlock()
		clientId := message["socketId"].(string)

		log.Debug("Resubscribe event for clientId: %v SocketId: %v", clientId, c.SocketId)
		found, err := c.Resubscribe(clientId)
		log.Debug("Resubscribe found for socketID: %v, %v", clientId, found)

		if err != nil {
			log.Error(err.Error())
		}
		sendToClient(c.Session, "broker.resubscribed", found)

	case "unsubscribe":
		globalMapMutex.Lock()
		defer globalMapMutex.Unlock()
		routingKeyPrefixes := message["routingKeyPrefix"].(string)
		log.Debug("Unsubscribe event for socketID: %v, and prefixes", c.SocketId, routingKeyPrefixes)
		for _, routingKeyPrefix := range strings.Split(routingKeyPrefixes, " ") {
			c.Unsubscribe(routingKeyPrefix)
		}

	case "publish":
		exchange := message["exchange"].(string)
		routingKey := message["routingKey"].(string)
		payload := message["payload"].(string)

		log.Debug("Publish Event: Exchange: %v, RoutingKey %v, Payload %v",
			exchange,
			routingKey,
			payload,
		)

		err := c.Publish(exchange, routingKey, payload)
		if err != nil {
			log.Error(err.Error())
		}

	case "ping":
		sendToClient(c.Session, "broker.pong", nil)

	default:
		log.Warning("Invalid action. message: %v socketId: %v", message, c.SocketId)

	}
}

// Publish publish the given payload for to the given exchange and routingkey.
func (c *Client) Publish(exchange, routingKey, payload string) error {
	if !strings.HasPrefix(routingKey, "client.") {
		return fmt.Errorf("Invalid routing key: %v socketId: %v", routingKey, c.SocketId)
	}

	for {
		c.LastPayload = ""
		err := c.ControlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: c.SocketId, Body: []byte(payload)})
		if err == nil {
			c.LastPayload = payload
			break
		}

		if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != amqp.ChannelError {
			log.Warning("payload: %v routing key: %v exchange: %v err: %v",
				payload, routingKey, exchange, err)
		}

		time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
		c.resetControlChannel()
	}

	return nil
}

// gaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// gaugeStart and calls the returned function in a defer statement.
func (c *Client) gaugeStart() (gaugeEnd func()) {
	log.Debug("Client connected: %v", c.Session.Tag)
	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if c.Session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", c.Session.Tag)
		changeClientsGauge(-1)
		if c.Session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

// resetControlChannel closes the current client's control channel and creates
// a new channel. It also listens to any server side error and publish back
// the error to the client.
func (c *Client) resetControlChannel() {
	if c.ControlChannel != nil {
		c.ControlChannel.Close()
	}

	var err error
	c.ControlChannel, err = c.Broker.PublishConn.Channel()
	if err != nil {
		panic(err)
	}

	go func() {
		defer log.RecoverAndLog()

		for amqpErr := range c.ControlChannel.NotifyClose(make(chan *amqp.Error)) {
			if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
				log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), c.LastPayload)
			}

			sendToClient(c.Session, "broker.error", map[string]interface{}{
				"code":    amqpErr.Code,
				"reason":  amqpErr.Reason,
				"server":  amqpErr.Server,
				"recover": amqpErr.Recover,
			})
		}
	}()
}

// RemoveFromRoute removes the sessions for the given routingKeyPrefix.
func (c *Client) RemoveFromRoute(routingKeyPrefix string) {
	routeSessions := routeMap[routingKeyPrefix]
	for i, routeSession := range routeSessions {
		if routeSession == c.Session {
			routeSessions[i] = routeSessions[len(routeSessions)-1]
			routeSessions = routeSessions[:len(routeSessions)-1]
			break
		}
	}
	if len(routeSessions) == 0 {
		delete(routeMap, routingKeyPrefix)
		return
	}
	routeMap[routingKeyPrefix] = routeSessions
}

// Subscribe add the given routingKeyPrefix to the list of subscriptions
// associated with this client.
func (c *Client) Subscribe(routingKeyPrefix string) error {
	res, err := c.Subscriptions.Has(routingKeyPrefix)
	if err != nil {
		return err
	}

	if res {
		return fmt.Errorf("Duplicate subscription to same routing key. %v %v", c.Session.Tag, routingKeyPrefix)
	}

	length, err := c.Subscriptions.Len()
	if err != nil {
		return err
	}

	if length > 0 && length%2000 == 0 {
		log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(length), c.Session.Tag)
	}

	globalMapMutex.Lock()
	routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], c.Session)
	globalMapMutex.Unlock()
	return c.Subscriptions.Subscribe(routingKeyPrefix)
}

func (c *Client) Resubscribe(sessionId string) (bool, error) {
	return c.Subscriptions.Resubscribe(sessionId)
}

// Unsubscribe deletes the given routingKey prefix from the subscription list
// and removes it from the global route map
func (c *Client) Unsubscribe(routingKeyPrefix string) {
	c.RemoveFromRoute(routingKeyPrefix)
	if err := c.Subscriptions.Unsubscribe(routingKeyPrefix); err != nil {
		fmt.Errorf("%v", err)
	}
}

// randomString() returns a new 16 char length random string
func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}
<commit_msg>Broker: add client's subscriptions to the routeMap<commit_after>package main

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"koding/broker/cache"
	"koding/tools/config"
	"koding/tools/sockjs"
	"strconv"
	"strings"
	"time"

	"github.com/streadway/amqp"
)

type Client struct {
	Session        *sockjs.Session
	ControlChannel *amqp.Channel
	SocketId       string
	Broker         *Broker
	LastPayload    string
	Subscriptions  *cache.SubscriptionStorage
}

// NewClient retuns a new client that is defined on a given session.
func NewClient(session *sockjs.Session, broker *Broker) *Client {
	socketID := randomString()
	session.Tag = socketID

	controlChannel, err := broker.PublishConn.Channel()
	if err != nil {
		panic(err)
	}

	subscriptions, err := cache.NewStorage(STORAGE_BACKEND, socketID)
	if err != nil {
		panic(err)
	}

	return &Client{
		Session:        session,
		SocketId:       socketID,
		ControlChannel: controlChannel,
		Broker:         broker,
		Subscriptions:  subscriptions,
	}
}

// Close should be called whenever a client disconnects.
func (c *Client) Close() {
	log.Debug("Client Close Request for socketID: %v", c.SocketId)
	c.Subscriptions.Each(func(routingKeyPrefix interface{}) bool {
		c.RemoveFromRoute(routingKeyPrefix.(string))
		return true
	})

	c.Subscriptions.ClearWithTimeout()

	for {
		err := c.ControlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(c.SocketId)})
		if err == nil {
			break
		}
		if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != amqp.ChannelError {
			panic(err)
		}
		c.resetControlChannel()
	}

	log.Debug("Closing control channel for socketID: %v", c.SocketId)
	c.ControlChannel.Close()
}

// handleSessionMessage handles the received message from the client. It
// passes a response back to the client or publish the received message to a
// rabbitmq exchange for further process.
func (c *Client) handleSessionMessage(data interface{}) {
	defer log.RecoverAndLog()

	message := data.(map[string]interface{})
	log.Debug("Received message: %v", message)

	action := message["action"]
	switch action {
	case "subscribe":
		for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
			if err := c.Subscribe(routingKeyPrefix); err != nil {
				log.Error(err.Error())
			}
		}

		sendToClient(c.Session, "broker.subscribed", message["routingKeyPrefix"])

	case "resubscribe":
		clientId := message["socketId"].(string)
		log.Debug("Resubscribe event for clientId: %v SocketId: %v", clientId, c.SocketId)
		found, err := c.Resubscribe(clientId)
		if err != nil {
			log.Error(err.Error())
		}
		log.Debug("Resubscribe found for socketID: %v, %v", clientId, found)
		sendToClient(c.Session, "broker.resubscribed", found)

	case "unsubscribe":
		routingKeyPrefixes := message["routingKeyPrefix"].(string)
		log.Debug("Unsubscribe event for socketID: %v, and prefixes", c.SocketId, routingKeyPrefixes)
		for _, routingKeyPrefix := range strings.Split(routingKeyPrefixes, " ") {
			c.Unsubscribe(routingKeyPrefix)
		}

	case "publish":
		exchange := message["exchange"].(string)
		routingKey := message["routingKey"].(string)
		payload := message["payload"].(string)

		log.Debug("Publish Event: Exchange: %v, RoutingKey %v, Payload %v",
			exchange,
			routingKey,
			payload,
		)

		err := c.Publish(exchange, routingKey, payload)
		if err != nil {
			log.Error(err.Error())
		}

	case "ping":
		sendToClient(c.Session, "broker.pong", nil)

	default:
		log.Warning("Invalid action. message: %v socketId: %v", message, c.SocketId)

	}
}

// Publish publish the given payload for to the given exchange and routingkey.
func (c *Client) Publish(exchange, routingKey, payload string) error {
	if !strings.HasPrefix(routingKey, "client.") {
		return fmt.Errorf("Invalid routing key: %v socketId: %v", routingKey, c.SocketId)
	}

	for {
		c.LastPayload = ""
		err := c.ControlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: c.SocketId, Body: []byte(payload)})
		if err == nil {
			c.LastPayload = payload
			break
		}

		if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != amqp.ChannelError {
			log.Warning("payload: %v routing key: %v exchange: %v err: %v",
				payload, routingKey, exchange, err)
		}

		time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
		c.resetControlChannel()
	}

	return nil
}

// gaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// gaugeStart and calls the returned function in a defer statement.
func (c *Client) gaugeStart() (gaugeEnd func()) {
	log.Debug("Client connected: %v", c.Session.Tag)
	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if c.Session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", c.Session.Tag)
		changeClientsGauge(-1)
		if c.Session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

// resetControlChannel closes the current client's control channel and creates
// a new channel. It also listens to any server side error and publish back
// the error to the client.
func (c *Client) resetControlChannel() {
	if c.ControlChannel != nil {
		c.ControlChannel.Close()
	}

	var err error
	c.ControlChannel, err = c.Broker.PublishConn.Channel()
	if err != nil {
		panic(err)
	}

	go func() {
		defer log.RecoverAndLog()

		for amqpErr := range c.ControlChannel.NotifyClose(make(chan *amqp.Error)) {
			if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
				log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), c.LastPayload)
			}

			sendToClient(c.Session, "broker.error", map[string]interface{}{
				"code":    amqpErr.Code,
				"reason":  amqpErr.Reason,
				"server":  amqpErr.Server,
				"recover": amqpErr.Recover,
			})
		}
	}()
}

// RemoveFromRoute removes the sessions for the given routingKeyPrefix.
func (c *Client) RemoveFromRoute(routingKeyPrefix string) {
	routeSessions := routeMap[routingKeyPrefix]
	for i, routeSession := range routeSessions {
		if routeSession == c.Session {
			routeSessions[i] = routeSessions[len(routeSessions)-1]
			routeSessions = routeSessions[:len(routeSessions)-1]
			break
		}
	}
	if len(routeSessions) == 0 {
		delete(routeMap, routingKeyPrefix)
		return
	}
	routeMap[routingKeyPrefix] = routeSessions
}

// Add to route
// todo ~ check for multiple subscriptions
func (c *Client) AddToRoute() {
	globalMapMutex.Lock()
	c.Subscriptions.Each(func(routingKeyPrefix interface{}) bool {
		rkp := routingKeyPrefix.(string)
		routeMap[rkp] = append(routeMap[rkp], c.Session)
		return true
	})
	globalMapMutex.Unlock()

}

// Subscribe add the given routingKeyPrefix to the list of subscriptions
// associated with this client.
func (c *Client) Subscribe(routingKeyPrefix string) error {
	res, err := c.Subscriptions.Has(routingKeyPrefix)
	if err != nil {
		return err
	}

	if res {
		return fmt.Errorf("Duplicate subscription to same routing key. %v %v", c.Session.Tag, routingKeyPrefix)
	}

	length, err := c.Subscriptions.Len()
	if err != nil {
		return err
	}

	if length > 0 && length%2000 == 0 {
		log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(length), c.Session.Tag)
	}

	globalMapMutex.Lock()
	routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], c.Session)
	globalMapMutex.Unlock()
	return c.Subscriptions.Subscribe(routingKeyPrefix)
}

func (c *Client) Resubscribe(sessionId string) (bool, error) {
	found, err := c.Subscriptions.Resubscribe(sessionId)
	if err != nil {
		return false, err
	}

	if !found {
		return false, nil
	}

	c.AddToRoute()
	return true, nil
}

// Unsubscribe deletes the given routingKey prefix from the subscription list
// and removes it from the global route map
func (c *Client) Unsubscribe(routingKeyPrefix string) {
	c.RemoveFromRoute(routingKeyPrefix)
	if err := c.Subscriptions.Unsubscribe(routingKeyPrefix); err != nil {
		fmt.Errorf("%v", err)
	}
}

// randomString() returns a new 16 char length random string
func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}
<|endoftext|>


<commit_before>package config

import (
	"net/url"
	"path"
	"path/filepath"
	"time"

	"koding/kites/kloud/utils/object"

	"github.com/boltdb/bolt"
	jwt "github.com/dgrijalva/jwt-go"
	konfig "github.com/koding/kite/config"
	"github.com/koding/kite/kitekey"
)

var KonfigCache = &CacheOptions{
	File: filepath.Join(KodingHome(), "konfig.bolt"),
	BoltDB: &bolt.Options{
		Timeout: 5 * time.Second,
	},
	Bucket: []byte("konfig"),
}

type Konfig struct {
	Environment string `json:"environment,omitempty"`

	// Kite configuration.
	KiteKeyFile string `json:"kiteKeyFile,omitempty"`
	KiteKey     string `json:"kiteKey,omitempty"`

	// Koding endpoints configuration.
	KontrolURL string `json:"kontrolURL,omitempty"`
	KlientURL  string `json:"klientURL,omitempty"`
	KloudURL   string `json:"kloudURL,omitempty"`
	TunnelURL  string `json:"tunnelURL,omitempty"`
	RemoteURL  string `json:"remoteURL,omitempty"`
	IPURL      string `json:"ipURL,omitempty"`
	IPCheckURL string `json:"ipCheckURL,omitempty"`

	// Koding networking configuration.
	//
	// TODO(rjeczalik): store command line flags in konfig.bolt
	// per Koding executable (KD / Klient).
	TunnelID string `json:"tunnelID,omitempty"`

	// Klient / KD auto-update endpoints.
	KlientLatestURL string `json:"klientLatestURL,omitempty"`
	KDLatestURL     string `json:"kdLatestURL,omitempty"`

	// Public S3 bucket for writing logs.
	PublicBucketName   string `json:"publicBucketName,omitempty"`
	PublicBucketRegion string `json:"publicBucketRegion,omitempty"`

	Debug bool `json:"debug,omitempty"`
}

func (k *Konfig) KiteHome() string {
	return filepath.Dir(k.KiteKeyFile)
}

func (k *Konfig) KiteConfig() *konfig.Config {
	return k.buildKiteConfig()
}

func (k *Konfig) KlientGzURL() string {
	u, err := url.Parse(k.KlientLatestURL)
	if err != nil {
		return ""
	}

	u.Path = path.Join(path.Dir(u.Path), "latest", "klient.gz")

	return u.String()
}

func (k *Konfig) buildKiteConfig() *konfig.Config {
	if k.KiteKey != "" {
		tok, err := jwt.ParseWithClaims(k.KiteKey, &kitekey.KiteClaims{}, kitekey.GetKontrolKey)
		if err == nil {
			cfg := &konfig.Config{}

			if err = cfg.ReadToken(tok); err == nil {
				return cfg
			}
		}
	}

	if k.KiteKeyFile != "" {
		if cfg, err := konfig.NewFromKiteKey(k.KiteKeyFile); err == nil {
			return cfg
		}
	}

	if cfg, err := konfig.Get(); err == nil {
		return cfg
	}

	return konfig.New()
}

// Enviroment is a hacky workaround for kd <-> klient environments.
// The managed klient expects to have kd from production channel,
// and devmanaged klient - from development. Depending from which
// app we load the the default Koding configuration, we need
// to cross-map the environments.
//
// TODO(rjeczalik): This should be fixed by removing managed / devmanaged
// channels - that is to improve "connect your vm" modal to
// not depend on special klient environment.
type Environments struct {
	Env       string
	KlientEnv string // Env is used if empty
	KDEnv     string // Env is used if empty
}

func (e *Environments) klientEnv() string {
	if e.KlientEnv != "" {
		return e.KlientEnv
	}
	return e.Env
}

func (e *Environments) kdEnv() string {
	if e.KDEnv != "" {
		return e.KDEnv
	}
	return e.Env
}

func NewKonfig(e *Environments) *Konfig {
	return &Konfig{
		Environment:        e.Env,
		KiteKeyFile:        "/etc/kite/kite.key",
		KlientURL:          "http://127.0.0.1:56789/kite",
		KontrolURL:         Builtin.Endpoints.Kontrol,
		KloudURL:           Builtin.Endpoints.Kloud,
		TunnelURL:          Builtin.Endpoints.TunnelServer,
		RemoteURL:          Builtin.Endpoints.RemoteAPI,
		IPURL:              Builtin.Endpoints.IP,
		IPCheckURL:         Builtin.Endpoints.IPCheck,
		KlientLatestURL:    ReplaceEnv(Builtin.Endpoints.KlientLatest, e.klientEnv()),
		KDLatestURL:        ReplaceEnv(Builtin.Endpoints.KDLatest, RmManaged(e.kdEnv())),
		PublicBucketName:   Builtin.Buckets.PublicLogs.Name,
		PublicBucketRegion: Builtin.Buckets.PublicLogs.Region,
		Debug:              false,
	}
}

func ReadKonfig(e *Environments) *Konfig {
	c := NewCache(KonfigCache)
	defer c.Close()

	return ReadKonfigFromCache(e, c)
}

func ReadKonfigFromCache(e *Environments, c *Cache) *Konfig {
	var override Konfig
	var builtin = NewKonfig(e)

	if err := c.GetValue("konfig", &override); err == nil {
		object.Merge(builtin, &override)
	}

	return builtin
}
<commit_msg>kd: add kodingBaseURL method<commit_after>package config

import (
	"net/url"
	"path"
	"path/filepath"
	"time"

	"koding/kites/kloud/utils/object"

	"github.com/boltdb/bolt"
	jwt "github.com/dgrijalva/jwt-go"
	konfig "github.com/koding/kite/config"
	"github.com/koding/kite/kitekey"
)

var KonfigCache = &CacheOptions{
	File: filepath.Join(KodingHome(), "konfig.bolt"),
	BoltDB: &bolt.Options{
		Timeout: 5 * time.Second,
	},
	Bucket: []byte("konfig"),
}

type Konfig struct {
	Environment string `json:"environment,omitempty"`

	// Kite configuration.
	KiteKeyFile string `json:"kiteKeyFile,omitempty"`
	KiteKey     string `json:"kiteKey,omitempty"`

	// Koding endpoints configuration.
	KontrolURL string `json:"kontrolURL,omitempty"`
	KlientURL  string `json:"klientURL,omitempty"`
	KloudURL   string `json:"kloudURL,omitempty"`
	KodingURL  string `json:"kodingURL,omitempty"`
	TunnelURL  string `json:"tunnelURL,omitempty"`
	RemoteURL  string `json:"remoteURL,omitempty"`
	IPURL      string `json:"ipURL,omitempty"`
	IPCheckURL string `json:"ipCheckURL,omitempty"`

	// Koding networking configuration.
	//
	// TODO(rjeczalik): store command line flags in konfig.bolt
	// per Koding executable (KD / Klient).
	TunnelID string `json:"tunnelID,omitempty"`

	// Klient / KD auto-update endpoints.
	KlientLatestURL string `json:"klientLatestURL,omitempty"`
	KDLatestURL     string `json:"kdLatestURL,omitempty"`

	// Public S3 bucket for writing logs.
	PublicBucketName   string `json:"publicBucketName,omitempty"`
	PublicBucketRegion string `json:"publicBucketRegion,omitempty"`

	Debug bool `json:"debug,omitempty"`
}

func (k *Konfig) KiteHome() string {
	return filepath.Dir(k.KiteKeyFile)
}

func (k *Konfig) KiteConfig() *konfig.Config {
	return k.buildKiteConfig()
}

func (k *Konfig) KlientGzURL() string {
	u, err := url.Parse(k.KlientLatestURL)
	if err != nil {
		return ""
	}

	u.Path = path.Join(path.Dir(u.Path), "latest", "klient.gz")

	return u.String()
}

func (k *Konfig) buildKiteConfig() *konfig.Config {
	if k.KiteKey != "" {
		tok, err := jwt.ParseWithClaims(k.KiteKey, &kitekey.KiteClaims{}, kitekey.GetKontrolKey)
		if err == nil {
			cfg := &konfig.Config{}

			if err = cfg.ReadToken(tok); err == nil {
				return cfg
			}
		}
	}

	if k.KiteKeyFile != "" {
		if cfg, err := konfig.NewFromKiteKey(k.KiteKeyFile); err == nil {
			return cfg
		}
	}

	if cfg, err := konfig.Get(); err == nil {
		return cfg
	}

	return konfig.New()
}

// Enviroment is a hacky workaround for kd <-> klient environments.
// The managed klient expects to have kd from production channel,
// and devmanaged klient - from development. Depending from which
// app we load the the default Koding configuration, we need
// to cross-map the environments.
//
// TODO(rjeczalik): This should be fixed by removing managed / devmanaged
// channels - that is to improve "connect your vm" modal to
// not depend on special klient environment.
type Environments struct {
	Env       string
	KlientEnv string // Env is used if empty
	KDEnv     string // Env is used if empty
}

func (e *Environments) klientEnv() string {
	if e.KlientEnv != "" {
		return e.KlientEnv
	}
	return e.Env
}

func (e *Environments) kdEnv() string {
	if e.KDEnv != "" {
		return e.KDEnv
	}
	return e.Env
}

func NewKonfig(e *Environments) *Konfig {
	return &Konfig{
		Environment:        e.Env,
		KiteKeyFile:        "/etc/kite/kite.key",
		KlientURL:          "http://127.0.0.1:56789/kite",
		KontrolURL:         Builtin.Endpoints.Kontrol,
		KloudURL:           Builtin.Endpoints.Kloud,
		TunnelURL:          Builtin.Endpoints.TunnelServer,
		RemoteURL:          Builtin.Endpoints.RemoteAPI,
		IPURL:              Builtin.Endpoints.IP,
		IPCheckURL:         Builtin.Endpoints.IPCheck,
		KlientLatestURL:    ReplaceEnv(Builtin.Endpoints.KlientLatest, e.klientEnv()),
		KDLatestURL:        ReplaceEnv(Builtin.Endpoints.KDLatest, RmManaged(e.kdEnv())),
		PublicBucketName:   Builtin.Buckets.PublicLogs.Name,
		PublicBucketRegion: Builtin.Buckets.PublicLogs.Region,
		Debug:              false,
	}
}

func ReadKonfig(e *Environments) *Konfig {
	c := NewCache(KonfigCache)
	defer c.Close()

	return ReadKonfigFromCache(e, c)
}

func ReadKonfigFromCache(e *Environments, c *Cache) *Konfig {
	var override Konfig
	var builtin = NewKonfig(e)

	if err := c.GetValue("konfig", &override); err == nil {
		object.Merge(builtin, &override)
	}

	return builtin
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"koding/kite-handler/command"
	"koding/kite-handler/fs"
	"koding/kite-handler/terminal"
	"koding/kites/klient/collaboration"
	"koding/kites/klient/protocol"
	"koding/kites/klient/usage"
	"log"
	"net"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/koding/kite"
	"github.com/koding/kite/config"
)

var (
	flagIP          = flag.String("ip", "", "Change public ip")
	flagPort        = flag.Int("port", 56789, "Change running port")
	flagVersion     = flag.Bool("version", false, "Show version and exit")
	flagProxy       = flag.Bool("proxy", false, "Start klient behind a proxy")
	flagEnvironment = flag.String("env", protocol.Environment, "Change environment")
	flagRegion      = flag.String("region", protocol.Region, "Change region")
	flagRegisterURL = flag.String("register-url", "", "Change register URL to kontrol")

	// update paramters
	flagUpdateInterval = flag.Duration("update-interval", time.Minute*5,
		"Change interval for checking for new updates")
	flagUpdateURL = flag.String("update-url",
		"https://s3.amazonaws.com/koding-klient/"+protocol.Environment+"/latest-version.txt",
		"Change update endpoint for latest version")

	VERSION = protocol.Version
	NAME    = protocol.Name

	// this is our main reference to count and measure metrics for the klient
	usg  = usage.NewUsage()
	klog kite.Logger

	collab = collaboration.New()

	// we also could use an atomic boolean this is simple for now.
	updating   = false
	updatingMu sync.Mutex // protects updating
)

func main() {
	flag.Parse()
	if *flagVersion {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	k := kite.New(NAME, VERSION)
	conf := config.MustGet()
	k.Config = conf
	k.Config.Port = *flagPort
	k.Config.Environment = *flagEnvironment
	k.Config.Region = *flagRegion

	// FIXME: It's ugly I know. It's a fix for Koding local development and is
	// needed
	if !strings.Contains(k.Config.KontrolURL, "ngrok") {
		// override current kontrolURL so it talks to port 3000, this is needed
		// because ELB can forward requests based on ports. The port 80 and 443 are
		// HTTP/HTTPS only so our kite can't connect it (we use websocket). However
		// We have a TCP proxy at 3000 which allows us to connect via WebSocket.
		u, _ := url.Parse(k.Config.KontrolURL)
		host, _, _ := net.SplitHostPort(u.Host)
		u.Host = AddPort(host, "3000")
		k.Config.KontrolURL = u.String()
	}

	klog = k.Log

	if *flagUpdateInterval < time.Minute {
		klog.Warning("Update interval can't be less than one minute. Setting to one minute.")
		*flagUpdateInterval = time.Minute
	}

	updater := &Updater{
		Endpoint: *flagUpdateURL,
		Interval: *flagUpdateInterval,
	}

	// before we register check for latest update and re-update itself before
	// we continue
	k.Log.Info("Checking for new updates")
	if err := updater.checkAndUpdate(); err != nil {
		klog.Warning("Self-update: %s", err)
	}

	go updater.Run()

	// always boot up with the same id in the kite.key
	k.Id = conf.Id

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

	// don't pass any request if the caller is outside of our scope.
	// don't allow anyone to call a method if we are during an update.
	k.PreHandleFunc(func(r *kite.Request) (interface{}, error) {
		// only authenticated methods have correct username. For example
		// kite.ping has authentication disabled so username can be empty.
		if r.Auth != nil {
			k.Log.Info("Kite '%s/%s/%s' called method: '%s'",
				r.Username, r.Client.Environment, r.Client.Name, r.Method)

			// Allow these users by default.
			allowedUsers := []string{k.Config.Username, "koding"}

			// Add collaboration users to the list
			for user := range collab.AllowedUsers {
				allowedUsers = append(allowedUsers, user)
			}

			if !userIn(r.Username, allowedUsers...) {
				return nil, fmt.Errorf("User '%s' is not allowed to make a call to us.", r.Username)
			}
		}

		updatingMu.Lock()
		defer updatingMu.Unlock()

		if updating {
			return nil, errors.New("Updating klient. Can't accept any method.")
		}

		return true, nil
	})

	// count only those methods, please add/remove methods here that will reset
	// the timer of a klient.
	usg.CountedMethods = map[string]bool{
		"fs.readDirectory":    true,
		"fs.glob":             true,
		"fs.readFile":         true,
		"fs.writeFile":        true,
		"fs.uniquePath":       true,
		"fs.getInfo":          true,
		"fs.setPermissions":   true,
		"fs.remove":           true,
		"fs.rename":           true,
		"fs.createDirectory":  true,
		"fs.move":             true,
		"fs.copy":             true,
		"webterm.getSessions": true,
		"webterm.connect":     true,
		"webterm.killSession": true,
		"exec":                true,
	}

	// we measure every incoming request
	k.PreHandleFunc(usg.Counter)

	// this provides us to get the current usage whenever we want
	k.HandleFunc("klient.usage", usg.Current)

	// also invoke updating
	k.Handle("klient.update", updater)

	// Collaboration
	k.HandleFunc("klient.share", collab.Share)
	k.HandleFunc("klient.unshare", collab.Unshare)
	k.HandleFunc("klient.shared", collab.Shared)

	k.HandleFunc("fs.readDirectory", fs.ReadDirectory)
	k.HandleFunc("fs.glob", fs.Glob)
	k.HandleFunc("fs.readFile", fs.ReadFile)
	k.HandleFunc("fs.writeFile", fs.WriteFile)
	k.HandleFunc("fs.uniquePath", fs.UniquePath)
	k.HandleFunc("fs.getInfo", fs.GetInfo)
	k.HandleFunc("fs.setPermissions", fs.SetPermissions)
	k.HandleFunc("fs.remove", fs.Remove)
	k.HandleFunc("fs.rename", fs.Rename)
	k.HandleFunc("fs.createDirectory", fs.CreateDirectory)
	k.HandleFunc("fs.move", fs.Move)
	k.HandleFunc("fs.copy", fs.Copy)

	terminal.ResetFunc = usg.Reset

	k.HandleFunc("webterm.getSessions", terminal.GetSessions)
	k.HandleFunc("webterm.connect", terminal.Connect)

	k.HandleFunc("webterm.killSession", terminal.KillSession)
	k.HandleFunc("exec", command.Exec)

	registerURL, err := getRegisterURL(k)
	if err != nil {
		log.Panic("could not get public ip" + err.Error())
	}

	if *flagRegisterURL != "" {
		u, err := url.Parse(*flagRegisterURL)
		if err != nil {
			k.Log.Fatal("Couldn't parse register url: %s", err)
		}

		registerURL = u
	}

	if registerURL == nil {
		log.Panic("register url is nil")
	}

	k.Log.Info("Going to register to kontrol with URL: %s", registerURL)
	if err := k.RegisterForever(registerURL); err != nil {
		log.Panic(err)
	}

	k.Log.Info("Running as version %s", VERSION)

	k.Run()
}

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func HasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func AddPort(host, port string) string {
	if ok := HasPort(host); ok {
		return host
	}

	return host + ":" + port
}
<commit_msg>klient: fix host, port splitting<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"koding/kite-handler/command"
	"koding/kite-handler/fs"
	"koding/kite-handler/terminal"
	"koding/kites/klient/collaboration"
	"koding/kites/klient/protocol"
	"koding/kites/klient/usage"
	"log"
	"net"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/koding/kite"
	"github.com/koding/kite/config"
)

var (
	flagIP          = flag.String("ip", "", "Change public ip")
	flagPort        = flag.Int("port", 56789, "Change running port")
	flagVersion     = flag.Bool("version", false, "Show version and exit")
	flagProxy       = flag.Bool("proxy", false, "Start klient behind a proxy")
	flagEnvironment = flag.String("env", protocol.Environment, "Change environment")
	flagRegion      = flag.String("region", protocol.Region, "Change region")
	flagRegisterURL = flag.String("register-url", "", "Change register URL to kontrol")

	// update paramters
	flagUpdateInterval = flag.Duration("update-interval", time.Minute*5,
		"Change interval for checking for new updates")
	flagUpdateURL = flag.String("update-url",
		"https://s3.amazonaws.com/koding-klient/"+protocol.Environment+"/latest-version.txt",
		"Change update endpoint for latest version")

	VERSION = protocol.Version
	NAME    = protocol.Name

	// this is our main reference to count and measure metrics for the klient
	usg  = usage.NewUsage()
	klog kite.Logger

	collab = collaboration.New()

	// we also could use an atomic boolean this is simple for now.
	updating   = false
	updatingMu sync.Mutex // protects updating
)

func main() {
	flag.Parse()
	if *flagVersion {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	k := kite.New(NAME, VERSION)
	conf := config.MustGet()
	k.Config = conf
	k.Config.Port = *flagPort
	k.Config.Environment = *flagEnvironment
	k.Config.Region = *flagRegion

	// FIXME: It's ugly I know. It's a fix for Koding local development and is
	// needed
	if !strings.Contains(k.Config.KontrolURL, "ngrok") {
		// override current kontrolURL so it talks to port 3000, this is needed
		// because ELB can forward requests based on ports. The port 80 and 443 are
		// HTTP/HTTPS only so our kite can't connect it (we use websocket). However
		// We have a TCP proxy at 3000 which allows us to connect via WebSocket.
		u, _ := url.Parse(k.Config.KontrolURL)

		host := u.Host
		if HasPort(u.Host) {
			host, _, _ = net.SplitHostPort(u.Host)
		}

		u.Host = AddPort(host, "3000")
		k.Config.KontrolURL = u.String()
	}

	klog = k.Log

	if *flagUpdateInterval < time.Minute {
		klog.Warning("Update interval can't be less than one minute. Setting to one minute.")
		*flagUpdateInterval = time.Minute
	}

	updater := &Updater{
		Endpoint: *flagUpdateURL,
		Interval: *flagUpdateInterval,
	}

	// before we register check for latest update and re-update itself before
	// we continue
	k.Log.Info("Checking for new updates")
	if err := updater.checkAndUpdate(); err != nil {
		klog.Warning("Self-update: %s", err)
	}

	go updater.Run()

	// always boot up with the same id in the kite.key
	k.Id = conf.Id

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

	// don't pass any request if the caller is outside of our scope.
	// don't allow anyone to call a method if we are during an update.
	k.PreHandleFunc(func(r *kite.Request) (interface{}, error) {
		// only authenticated methods have correct username. For example
		// kite.ping has authentication disabled so username can be empty.
		if r.Auth != nil {
			k.Log.Info("Kite '%s/%s/%s' called method: '%s'",
				r.Username, r.Client.Environment, r.Client.Name, r.Method)

			// Allow these users by default.
			allowedUsers := []string{k.Config.Username, "koding"}

			// Add collaboration users to the list
			for user := range collab.AllowedUsers {
				allowedUsers = append(allowedUsers, user)
			}

			if !userIn(r.Username, allowedUsers...) {
				return nil, fmt.Errorf("User '%s' is not allowed to make a call to us.", r.Username)
			}
		}

		updatingMu.Lock()
		defer updatingMu.Unlock()

		if updating {
			return nil, errors.New("Updating klient. Can't accept any method.")
		}

		return true, nil
	})

	// count only those methods, please add/remove methods here that will reset
	// the timer of a klient.
	usg.CountedMethods = map[string]bool{
		"fs.readDirectory":    true,
		"fs.glob":             true,
		"fs.readFile":         true,
		"fs.writeFile":        true,
		"fs.uniquePath":       true,
		"fs.getInfo":          true,
		"fs.setPermissions":   true,
		"fs.remove":           true,
		"fs.rename":           true,
		"fs.createDirectory":  true,
		"fs.move":             true,
		"fs.copy":             true,
		"webterm.getSessions": true,
		"webterm.connect":     true,
		"webterm.killSession": true,
		"exec":                true,
	}

	// we measure every incoming request
	k.PreHandleFunc(usg.Counter)

	// this provides us to get the current usage whenever we want
	k.HandleFunc("klient.usage", usg.Current)

	// also invoke updating
	k.Handle("klient.update", updater)

	// Collaboration
	k.HandleFunc("klient.share", collab.Share)
	k.HandleFunc("klient.unshare", collab.Unshare)
	k.HandleFunc("klient.shared", collab.Shared)

	k.HandleFunc("fs.readDirectory", fs.ReadDirectory)
	k.HandleFunc("fs.glob", fs.Glob)
	k.HandleFunc("fs.readFile", fs.ReadFile)
	k.HandleFunc("fs.writeFile", fs.WriteFile)
	k.HandleFunc("fs.uniquePath", fs.UniquePath)
	k.HandleFunc("fs.getInfo", fs.GetInfo)
	k.HandleFunc("fs.setPermissions", fs.SetPermissions)
	k.HandleFunc("fs.remove", fs.Remove)
	k.HandleFunc("fs.rename", fs.Rename)
	k.HandleFunc("fs.createDirectory", fs.CreateDirectory)
	k.HandleFunc("fs.move", fs.Move)
	k.HandleFunc("fs.copy", fs.Copy)

	terminal.ResetFunc = usg.Reset

	k.HandleFunc("webterm.getSessions", terminal.GetSessions)
	k.HandleFunc("webterm.connect", terminal.Connect)

	k.HandleFunc("webterm.killSession", terminal.KillSession)
	k.HandleFunc("exec", command.Exec)

	registerURL, err := getRegisterURL(k)
	if err != nil {
		log.Panic("could not get public ip" + err.Error())
	}

	if *flagRegisterURL != "" {
		u, err := url.Parse(*flagRegisterURL)
		if err != nil {
			k.Log.Fatal("Couldn't parse register url: %s", err)
		}

		registerURL = u
	}

	if registerURL == nil {
		log.Panic("register url is nil")
	}

	k.Log.Info("Going to register to kontrol with URL: %s", registerURL)
	if err := k.RegisterForever(registerURL); err != nil {
		log.Panic(err)
	}

	k.Log.Info("Running as version %s", VERSION)

	k.Run()
}

// Given a string of the form "host", "host:port", or "[ipv6::address]:port",
// return true if the string includes a port.
func HasPort(s string) bool { return strings.LastIndex(s, ":") > strings.LastIndex(s, "]") }

// Given a string of the form "host", "port", returns "host:port"
func AddPort(host, port string) string {
	if ok := HasPort(host); ok {
		return host
	}

	return host + ":" + port
}
<|endoftext|>


<commit_before>package kite

import (
	"errors"
	"fmt"
	"github.com/op/go-logging"
	"koding/newkite/dnode"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"strconv"
	"sync"
	"time"
)

const DefaultTellTimeout = 4 * time.Second

// RemoteKite is the client for communicating with another Kite.
// It has Tell() and Go() methods for calling methods sync/async way.
type RemoteKite struct {
	// The information about the kite that we are connecting to.
	protocol.Kite

	// A reference to the current Kite running.
	localKite *Kite

	// A reference to the Kite's logger for easy access.
	Log *logging.Logger

	// Credentials that we sent in each request.
	Authentication Authentication

	// dnode RPC client that processes messages.
	client *rpc.Client

	// To signal waiters of Go() on disconnect.
	disconnect chan bool

	// Duration to wait reply from remote when making a request with Tell().
	tellTimeout time.Duration
}

// NewRemoteKite returns a pointer to a new RemoteKite. The returned instance
// is not connected. You have to call Dial() or DialForever() before calling
// Tell() and Go() methods.
func (k *Kite) NewRemoteKite(kite protocol.Kite, auth Authentication) *RemoteKite {
	r := &RemoteKite{
		Kite:           kite,
		localKite:      k,
		Log:            k.Log,
		Authentication: auth,
		client:         k.server.NewClientWithHandlers(),
		disconnect:     make(chan bool),
	}
	r.SetTellTimeout(DefaultTellTimeout)

	// Required for customizing dnode protocol for Kite.
	r.client.SetWrappers(wrapMethodArgs, wrapCallbackArgs, runMethod, runCallback)

	// We need a reference to the local kite when a method call is received.
	r.client.Properties()["localKite"] = k

	// We need a reference to the remote kite when sending a message to remote.
	r.client.Properties()["remoteKite"] = r

	r.OnConnect(func() {
		if r.Authentication.validUntil == nil {
			return
		}

		// Start a goroutine that will renew the token before it expires.
		go r.tokenRenewer()
	})

	var m sync.Mutex
	r.OnDisconnect(func() {
		m.Lock()
		close(r.disconnect)
		r.disconnect = make(chan bool)
		m.Unlock()
	})

	return r
}

func wrapCallbackArgs(args []interface{}, tr dnode.Transport) []interface{} {
	return []interface{}{&callOptionsOut{
		WithArgs: args,
		callOptions: callOptions{
			Kite: tr.Properties()["localKite"].(*Kite).Kite,
		},
	}}
}

// newRemoteKiteWithClient returns a pointer to new RemoteKite instance.
// The client will be replaced with the given client.
// Used to give the Kite method handler a working RemoteKite to call methods
// on other side.
func (k *Kite) newRemoteKiteWithClient(kite protocol.Kite, auth Authentication, client *rpc.Client) *RemoteKite {
	r := k.NewRemoteKite(kite, auth)
	r.client = client
	r.client.SetWrappers(wrapMethodArgs, wrapCallbackArgs, runMethod, runCallback)
	r.client.Properties()["localKite"] = k
	r.client.Properties()["remoteKite"] = r
	return r
}

// SetTellTimeout sets the timeout duration for requests made with Tell().
func (r *RemoteKite) SetTellTimeout(d time.Duration) { r.tellTimeout = d }

// Dial connects to the remote Kite. Returns error if it can't.
func (r *RemoteKite) Dial() (err error) {
	addr := r.Kite.Addr()
	r.Log.Info("Dialing remote kite: [%s %s]", r.Kite.Name, addr)
	return r.client.Dial("ws://" + addr + "/dnode")
}

// Dial connects to the remote Kite. If it can't connect, it retries indefinitely.
func (r *RemoteKite) DialForever() {
	addr := r.Kite.Addr()
	r.Log.Info("Dialing remote kite: [%s %s]", r.Kite.Name, addr)
	r.client.DialForever("ws://" + addr + "/dnode")
}

func (r *RemoteKite) Close() {
	r.client.Close()
}

// OnConnect registers a function to run on connect.
func (r *RemoteKite) OnConnect(handler func()) {
	r.client.OnConnect(handler)
}

// OnDisconnect registers a function to run on disconnect.
func (r *RemoteKite) OnDisconnect(handler func()) {
	r.client.OnDisconnect(handler)
}

func (r *RemoteKite) tokenRenewer() {
	for {
		// Token will be renewed before it expires.
		renewTime := r.Authentication.validUntil.Add(-30 * time.Second)
		select {
		case <-time.After(renewTime.Sub(time.Now().UTC())):
			if err := r.renewTokenUntilDisconnect(); err != nil {
				return
			}
		case <-r.disconnect:
			return
		}
	}
}

// renewToken retries until the request is successful or disconnect.
func (r *RemoteKite) renewTokenUntilDisconnect() error {
	const retryInterval = 10 * time.Second

	if err := r.renewToken(); err == nil {
		return nil
	}

loop:
	for {
		select {
		case <-time.After(retryInterval):
			if err := r.renewToken(); err != nil {
				r.Log.Error("error: %s Cannot renew token for Kite: %s I will retry in %d seconds...", err.Error(), r.Kite.ID, retryInterval)
				continue
			}

			break loop
		case <-r.disconnect:
			return errors.New("disconnect")
		}
	}

	return nil
}

func (r *RemoteKite) renewToken() error {
	tkn, err := r.localKite.Kontrol.GetToken(&r.Kite)
	if err != nil {
		return err
	}

	validUntil := time.Now().UTC().Add(time.Duration(tkn.TTL) * time.Second)
	r.Authentication.Key = tkn.Key
	r.Authentication.validUntil = &validUntil

	return nil
}

// callOptions is the type of first argument in the dnode message.
// Second argument is a callback function.
// It is used when unmarshalling a dnode message.
type callOptions struct {
	// Arguments to the method
	Kite             protocol.Kite    `json:"kite"`
	Authentication   Authentication   `json:"authentication"`
	WithArgs         []*dnode.Partial `json:"withArgs" dnode:"-"`
	ResponseCallback dnode.Function   `json:"responseCallback" dnode:"-"`
}

// callOptionsOut is the same structure with callOptions.
// It is used when marshalling a dnode message.
type callOptionsOut struct {
	callOptions

	// Override this when sending because args will not be a *dnode.Partial.
	WithArgs []interface{} `json:"withArgs"`

	// Override for sending. Incoming type is dnode.Function.
	ResponseCallback Callback `json:"responseCallback"`
}

// That's what we send as a first argument in dnode message.
func wrapMethodArgs(args []interface{}, tr dnode.Transport) []interface{} {
	r := tr.Properties()["remoteKite"].(*RemoteKite)

	responseCallback := args[len(args)-1:][0].(Callback)
	args = args[:len(args)-1]

	options := callOptionsOut{
		WithArgs:         args,
		ResponseCallback: responseCallback,
		callOptions: callOptions{
			Kite:           r.localKite.Kite,
			Authentication: r.Authentication,
		},
	}

	return []interface{}{options}
}

type Authentication struct {
	// Type can be "kodingKey", "token" or "sessionID" for now.
	Type       string     `json:"type"`
	Key        string     `json:"key"`
	validUntil *time.Time `json:"-"`
}

type response struct {
	Result *dnode.Partial
	Err    error
}

// Tell makes a blocking method call to the server.
// Waits until the callback function is called by the other side and
// returns the result and the error.
func (r *RemoteKite) Tell(method string, args ...interface{}) (result *dnode.Partial, err error) {
	return r.TellWithTimeout(method, 0, args...)
}

// TellWithTimeout does the same thing with Tell() method except it takes an
// extra argument that is the timeout for waiting reply from the remote Kite.
// If timeout is given 0, the behavior is same as Tell().
func (r *RemoteKite) TellWithTimeout(method string, timeout time.Duration, args ...interface{}) (result *dnode.Partial, err error) {
	response := <-r.GoWithTimeout(method, timeout, args...)
	return response.Result, response.Err
}

// Go makes an unblocking method call to the server.
// It returns a channel that the caller can wait on it to get the response.
func (r *RemoteKite) Go(method string, args ...interface{}) chan *response {
	return r.GoWithTimeout(method, 0, args...)
}

// GoWithTimeout does the same thing with Go() method except it takes an
// extra argument that is the timeout for waiting reply from the remote Kite.
// If timeout is given 0, the behavior is same as Go().
func (r *RemoteKite) GoWithTimeout(method string, timeout time.Duration, args ...interface{}) chan *response {
	// We will return this channel to the caller.
	// It can wait on this channel to get the response.
	r.Log.Debug("Telling method [%s] on kite [%s]", method, r.Name)
	responseChan := make(chan *response, 1)

	r.send(method, args, timeout, responseChan)

	return responseChan
}

// send sends the method with callback to the server.
func (r *RemoteKite) send(method string, args []interface{}, timeout time.Duration, responseChan chan *response) {
	// To clean the sent callback after response is received.
	// Send/Receive in a channel to prevent race condition because
	// the callback is run in a separate goroutine.
	removeCallback := make(chan uint64, 1)

	// When a callback is called it will send the response to this channel.
	doneChan := make(chan *response, 1)

	cb := r.makeResponseCallback(doneChan, removeCallback)
	args = append(args, cb)

	// BUG: This sometimes does not return an error, even if the remote
	// kite is disconnected. I could not find out why.
	// Timeout below in goroutine saves us in this case.
	callbacks, err := r.client.Call(method, args...)
	if err != nil {
		responseChan <- &response{
			Result: nil,
			Err: fmt.Errorf("Telling method [%s] on [%s] error: %s",
				method, r.Kite.Name, err),
		}
		return
	}

	// Use default timeout from r (RemoteKite) if zero.
	if timeout == 0 {
		timeout = r.tellTimeout
	}

	// Waits until the response has came or the connection has disconnected.
	go func() {
		select {
		case resp := <-doneChan:
			responseChan <- resp
		case <-r.disconnect:
			responseChan <- &response{nil, errors.New("Client disconnected")}
		case <-time.After(timeout):
			responseChan <- &response{nil, errors.New("Timeout")}

			// Remove the callback function from the map so we do not
			// consume memory for unused callbacks.
			if id, ok := <-removeCallback; ok {
				r.client.RemoveCallback(id)
			}
		}
	}()

	sendCallbackID(callbacks, removeCallback)
}

// sendCallbackID send the callback number to be deleted after response is received.
func sendCallbackID(callbacks map[string]dnode.Path, ch chan uint64) {
	if len(callbacks) > 0 {
		// Find max callback ID.
		max := uint64(0)
		for id, _ := range callbacks {
			i, _ := strconv.ParseUint(id, 10, 64)
			if i > max {
				max = i
			}
		}

		ch <- max
	} else {
		close(ch)
	}
}

// makeResponseCallback prepares and returns a callback function sent to the server.
// The caller of the Tell() is blocked until the server calls this callback function.
// Sets theResponse and notifies the caller by sending to done channel.
func (r *RemoteKite) makeResponseCallback(doneChan chan *response, removeCallback <-chan uint64) Callback {
	return Callback(func(request *Request) {
		var (
			err    error          // First argument
			result *dnode.Partial // Second argument
		)

		// Notify that the callback is finished.
		defer func() { doneChan <- &response{result, err} }()

		// Remove the callback function from the map so we do not
		// consume memory for unused callbacks.
		if id, ok := <-removeCallback; ok {
			r.client.RemoveCallback(id)
		}

		// Unmarshal callback response.
		responseArgs := request.Args.MustSliceOfLength(1)
		var response struct {
			Error  *Error         `json:"error"`
			Result *dnode.Partial `json:"result"`
		}
		responseArgs[0].MustUnmarshal(&response)

		result = response.Result

		// This is the error argument. Unmarshal panics if it is null.
		if response.Error == nil {
			return
		}

		// Read the error argument in response.
		var kiteErr Error
		responseArgs[0].MustUnmarshal(&kiteErr)

		r.Log.Warning("Error in remote Kite: %s", kiteErr.Error())
		err = &kiteErr
	})
}
<commit_msg>kite: add comment<commit_after>package kite

import (
	"errors"
	"fmt"
	"github.com/op/go-logging"
	"koding/newkite/dnode"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"strconv"
	"sync"
	"time"
)

const DefaultTellTimeout = 4 * time.Second

// RemoteKite is the client for communicating with another Kite.
// It has Tell() and Go() methods for calling methods sync/async way.
type RemoteKite struct {
	// The information about the kite that we are connecting to.
	protocol.Kite

	// A reference to the current Kite running.
	localKite *Kite

	// A reference to the Kite's logger for easy access.
	Log *logging.Logger

	// Credentials that we sent in each request.
	Authentication Authentication

	// dnode RPC client that processes messages.
	client *rpc.Client

	// To signal waiters of Go() on disconnect.
	disconnect chan bool

	// Duration to wait reply from remote when making a request with Tell().
	tellTimeout time.Duration
}

// NewRemoteKite returns a pointer to a new RemoteKite. The returned instance
// is not connected. You have to call Dial() or DialForever() before calling
// Tell() and Go() methods.
func (k *Kite) NewRemoteKite(kite protocol.Kite, auth Authentication) *RemoteKite {
	r := &RemoteKite{
		Kite:           kite,
		localKite:      k,
		Log:            k.Log,
		Authentication: auth,
		client:         k.server.NewClientWithHandlers(),
		disconnect:     make(chan bool),
	}
	r.SetTellTimeout(DefaultTellTimeout)

	// Required for customizing dnode protocol for Kite.
	r.client.SetWrappers(wrapMethodArgs, wrapCallbackArgs, runMethod, runCallback)

	// We need a reference to the local kite when a method call is received.
	r.client.Properties()["localKite"] = k

	// We need a reference to the remote kite when sending a message to remote.
	r.client.Properties()["remoteKite"] = r

	r.OnConnect(func() {
		if r.Authentication.validUntil == nil {
			return
		}

		// Start a goroutine that will renew the token before it expires.
		go r.tokenRenewer()
	})

	var m sync.Mutex
	r.OnDisconnect(func() {
		m.Lock()
		close(r.disconnect)
		r.disconnect = make(chan bool)
		m.Unlock()
	})

	return r
}

func wrapCallbackArgs(args []interface{}, tr dnode.Transport) []interface{} {
	return []interface{}{&callOptionsOut{
		WithArgs: args,
		callOptions: callOptions{
			Kite: tr.Properties()["localKite"].(*Kite).Kite,
		},
	}}
}

// newRemoteKiteWithClient returns a pointer to new RemoteKite instance.
// The client will be replaced with the given client.
// Used to give the Kite method handler a working RemoteKite to call methods
// on other side.
func (k *Kite) newRemoteKiteWithClient(kite protocol.Kite, auth Authentication, client *rpc.Client) *RemoteKite {
	r := k.NewRemoteKite(kite, auth)
	r.client = client
	r.client.SetWrappers(wrapMethodArgs, wrapCallbackArgs, runMethod, runCallback)
	r.client.Properties()["localKite"] = k
	r.client.Properties()["remoteKite"] = r
	return r
}

// SetTellTimeout sets the timeout duration for requests made with Tell().
func (r *RemoteKite) SetTellTimeout(d time.Duration) { r.tellTimeout = d }

// Dial connects to the remote Kite. Returns error if it can't.
func (r *RemoteKite) Dial() (err error) {
	addr := r.Kite.Addr()
	r.Log.Info("Dialing remote kite: [%s %s]", r.Kite.Name, addr)
	return r.client.Dial("ws://" + addr + "/dnode")
}

// Dial connects to the remote Kite. If it can't connect, it retries indefinitely.
func (r *RemoteKite) DialForever() {
	addr := r.Kite.Addr()
	r.Log.Info("Dialing remote kite: [%s %s]", r.Kite.Name, addr)
	r.client.DialForever("ws://" + addr + "/dnode")
}

func (r *RemoteKite) Close() {
	r.client.Close()
}

// OnConnect registers a function to run on connect.
func (r *RemoteKite) OnConnect(handler func()) {
	r.client.OnConnect(handler)
}

// OnDisconnect registers a function to run on disconnect.
func (r *RemoteKite) OnDisconnect(handler func()) {
	r.client.OnDisconnect(handler)
}

func (r *RemoteKite) tokenRenewer() {
	for {
		// Token will be renewed before it expires.
		renewTime := r.Authentication.validUntil.Add(-30 * time.Second)
		select {
		case <-time.After(renewTime.Sub(time.Now().UTC())):
			if err := r.renewTokenUntilDisconnect(); err != nil {
				return
			}
		case <-r.disconnect:
			return
		}
	}
}

// renewToken retries until the request is successful or disconnect.
func (r *RemoteKite) renewTokenUntilDisconnect() error {
	const retryInterval = 10 * time.Second

	if err := r.renewToken(); err == nil {
		return nil
	}

loop:
	for {
		select {
		case <-time.After(retryInterval):
			if err := r.renewToken(); err != nil {
				r.Log.Error("error: %s Cannot renew token for Kite: %s I will retry in %d seconds...", err.Error(), r.Kite.ID, retryInterval)
				continue
			}

			break loop
		case <-r.disconnect:
			return errors.New("disconnect")
		}
	}

	return nil
}

func (r *RemoteKite) renewToken() error {
	tkn, err := r.localKite.Kontrol.GetToken(&r.Kite)
	if err != nil {
		return err
	}

	validUntil := time.Now().UTC().Add(time.Duration(tkn.TTL) * time.Second)
	r.Authentication.Key = tkn.Key
	r.Authentication.validUntil = &validUntil

	return nil
}

// callOptions is the type of first argument in the dnode message.
// Second argument is a callback function.
// It is used when unmarshalling a dnode message.
type callOptions struct {
	// Arguments to the method
	Kite             protocol.Kite    `json:"kite"`
	Authentication   Authentication   `json:"authentication"`
	WithArgs         []*dnode.Partial `json:"withArgs" dnode:"-"`
	ResponseCallback dnode.Function   `json:"responseCallback" dnode:"-"`
}

// callOptionsOut is the same structure with callOptions.
// It is used when marshalling a dnode message.
type callOptionsOut struct {
	callOptions

	// Override this when sending because args will not be a *dnode.Partial.
	WithArgs []interface{} `json:"withArgs"`

	// Override for sending. Incoming type is dnode.Function.
	ResponseCallback Callback `json:"responseCallback"`
}

// That's what we send as a first argument in dnode message.
func wrapMethodArgs(args []interface{}, tr dnode.Transport) []interface{} {
	r := tr.Properties()["remoteKite"].(*RemoteKite)

	responseCallback := args[len(args)-1].(Callback) // last item
	args = args[:len(args)-1]                        // previout items

	options := callOptionsOut{
		WithArgs:         args,
		ResponseCallback: responseCallback,
		callOptions: callOptions{
			Kite:           r.localKite.Kite,
			Authentication: r.Authentication,
		},
	}

	return []interface{}{options}
}

type Authentication struct {
	// Type can be "kodingKey", "token" or "sessionID" for now.
	Type       string     `json:"type"`
	Key        string     `json:"key"`
	validUntil *time.Time `json:"-"`
}

type response struct {
	Result *dnode.Partial
	Err    error
}

// Tell makes a blocking method call to the server.
// Waits until the callback function is called by the other side and
// returns the result and the error.
func (r *RemoteKite) Tell(method string, args ...interface{}) (result *dnode.Partial, err error) {
	return r.TellWithTimeout(method, 0, args...)
}

// TellWithTimeout does the same thing with Tell() method except it takes an
// extra argument that is the timeout for waiting reply from the remote Kite.
// If timeout is given 0, the behavior is same as Tell().
func (r *RemoteKite) TellWithTimeout(method string, timeout time.Duration, args ...interface{}) (result *dnode.Partial, err error) {
	response := <-r.GoWithTimeout(method, timeout, args...)
	return response.Result, response.Err
}

// Go makes an unblocking method call to the server.
// It returns a channel that the caller can wait on it to get the response.
func (r *RemoteKite) Go(method string, args ...interface{}) chan *response {
	return r.GoWithTimeout(method, 0, args...)
}

// GoWithTimeout does the same thing with Go() method except it takes an
// extra argument that is the timeout for waiting reply from the remote Kite.
// If timeout is given 0, the behavior is same as Go().
func (r *RemoteKite) GoWithTimeout(method string, timeout time.Duration, args ...interface{}) chan *response {
	// We will return this channel to the caller.
	// It can wait on this channel to get the response.
	r.Log.Debug("Telling method [%s] on kite [%s]", method, r.Name)
	responseChan := make(chan *response, 1)

	r.send(method, args, timeout, responseChan)

	return responseChan
}

// send sends the method with callback to the server.
func (r *RemoteKite) send(method string, args []interface{}, timeout time.Duration, responseChan chan *response) {
	// To clean the sent callback after response is received.
	// Send/Receive in a channel to prevent race condition because
	// the callback is run in a separate goroutine.
	removeCallback := make(chan uint64, 1)

	// When a callback is called it will send the response to this channel.
	doneChan := make(chan *response, 1)

	cb := r.makeResponseCallback(doneChan, removeCallback)
	args = append(args, cb)

	// BUG: This sometimes does not return an error, even if the remote
	// kite is disconnected. I could not find out why.
	// Timeout below in goroutine saves us in this case.
	callbacks, err := r.client.Call(method, args...)
	if err != nil {
		responseChan <- &response{
			Result: nil,
			Err: fmt.Errorf("Telling method [%s] on [%s] error: %s",
				method, r.Kite.Name, err),
		}
		return
	}

	// Use default timeout from r (RemoteKite) if zero.
	if timeout == 0 {
		timeout = r.tellTimeout
	}

	// Waits until the response has came or the connection has disconnected.
	go func() {
		select {
		case resp := <-doneChan:
			responseChan <- resp
		case <-r.disconnect:
			responseChan <- &response{nil, errors.New("Client disconnected")}
		case <-time.After(timeout):
			responseChan <- &response{nil, errors.New("Timeout")}

			// Remove the callback function from the map so we do not
			// consume memory for unused callbacks.
			if id, ok := <-removeCallback; ok {
				r.client.RemoveCallback(id)
			}
		}
	}()

	sendCallbackID(callbacks, removeCallback)
}

// sendCallbackID send the callback number to be deleted after response is received.
func sendCallbackID(callbacks map[string]dnode.Path, ch chan uint64) {
	if len(callbacks) > 0 {
		// Find max callback ID.
		max := uint64(0)
		for id, _ := range callbacks {
			i, _ := strconv.ParseUint(id, 10, 64)
			if i > max {
				max = i
			}
		}

		ch <- max
	} else {
		close(ch)
	}
}

// makeResponseCallback prepares and returns a callback function sent to the server.
// The caller of the Tell() is blocked until the server calls this callback function.
// Sets theResponse and notifies the caller by sending to done channel.
func (r *RemoteKite) makeResponseCallback(doneChan chan *response, removeCallback <-chan uint64) Callback {
	return Callback(func(request *Request) {
		var (
			err    error          // First argument
			result *dnode.Partial // Second argument
		)

		// Notify that the callback is finished.
		defer func() { doneChan <- &response{result, err} }()

		// Remove the callback function from the map so we do not
		// consume memory for unused callbacks.
		if id, ok := <-removeCallback; ok {
			r.client.RemoveCallback(id)
		}

		// Unmarshal callback response.
		responseArgs := request.Args.MustSliceOfLength(1)
		var response struct {
			Error  *Error         `json:"error"`
			Result *dnode.Partial `json:"result"`
		}
		responseArgs[0].MustUnmarshal(&response)

		result = response.Result

		// This is the error argument. Unmarshal panics if it is null.
		if response.Error == nil {
			return
		}

		// Read the error argument in response.
		var kiteErr Error
		responseArgs[0].MustUnmarshal(&kiteErr)

		r.Log.Warning("Error in remote Kite: %s", kiteErr.Error())
		err = &kiteErr
	})
}
<|endoftext|>


<commit_before>package config

import (
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/exec"
)

type Config struct {
	BuildNumber     int
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
	}
	Mongo string
	Mq    struct {
		Host          string
		Port          int
		ComponentUser string
		Password      string
		Vhost         string
	}
	Neo4j struct {
		Read    string
		Write   string
		Port    int
		Enabled bool
	}
	Broker struct {
		IP              string
		Port            int
		CertFile        string
		KeyFile         string
		AuthExchange    string
		AuthAllExchange string
	}
	Loggr struct {
		Push   bool
		Url    string
		ApiKey string
	}
	Librato struct {
		Push     bool
		Email    string
		Token    string
		Interval int
	}
	Opsview struct {
		Push bool
		Host string
	}
	Kontrold struct {
		Api struct {
			Port int
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
			SSLIPS  string
		}
		RabbitMq struct {
			Host     string
			Port     string
			Login    string
			Password string
			Vhost    string
		}
	}
	FollowFeed struct {
		Host          string
		Port          int
		ComponentUser string
		Password      string
		Vhost         string
	}
}

var FileProfile string
var PillarProfile string
var Profile string
var Current Config
var LogDebug bool
var Uuid string
var Host string
var BrokerDomain string
var Region string

func init() {
	flag.StringVar(&FileProfile, "c", "", "Configuration profile from file")
	flag.StringVar(&FileProfile, "config", "", "Alias for -c")
	flag.StringVar(&PillarProfile, "p", "", "Configuration profile from saltstack pillar")
	flag.BoolVar(&LogDebug, "d", false, "Log debug messages")
	flag.StringVar(&Uuid, "u", "", "Enable kontrol mode")
	flag.StringVar(&Host, "h", "", "Hostname to be resolved")
	flag.StringVar(&BrokerDomain, "a", "", "Send kontrol a custom domain istead of os.Hostname")
	flag.StringVar(&BrokerDomain, "domain", "", "Alias for -a")
	flag.StringVar(&Region, "r", "", "Region")

	flag.Parse()
	if flag.NArg() != 0 {
		flag.PrintDefaults()
		os.Exit(1)
	}
	if FileProfile == "" && PillarProfile == "" {
		fmt.Println("Please specify a configuration profile via -c or -p.")
		flag.PrintDefaults()
		os.Exit(1)
	}
	if FileProfile != "" && PillarProfile != "" {
		fmt.Println("The flags -c and -p are exclusive.")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var configCommand *exec.Cmd
	if FileProfile != "" {
		Profile = FileProfile
		configCommand = exec.Command("node", "-e", "require('koding-config-manager').printJson('main."+FileProfile+"')")
	}
	if PillarProfile != "" {
		Profile = PillarProfile
		configCommand = exec.Command("salt-call", "pillar.get", PillarProfile, "--output=json", "--log-level=warning")
	}

	configJSON, err := configCommand.CombinedOutput()
	if err != nil {
		fmt.Printf("Could not execute configuration source: %s\nConfiguration source output:\n%s\n", err.Error(), configJSON)
		os.Exit(1)
	}

	err = json.Unmarshal(configJSON, &Current)
	if err == nil && (Current == Config{}) {
		err = fmt.Errorf("Empty configuration.")
	}
	if err != nil {
		fmt.Printf("Could not unmarshal configuration: %s\nConfiguration source output:\n%s\n", err.Error(), configJSON)
		os.Exit(1)
	}
}
<commit_msg>go: Toggle config.LogDebug with SIGUSR2.<commit_after>package config

import (
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
)

type Config struct {
	BuildNumber     int
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
	}
	Mongo string
	Mq    struct {
		Host          string
		Port          int
		ComponentUser string
		Password      string
		Vhost         string
	}
	Neo4j struct {
		Read    string
		Write   string
		Port    int
		Enabled bool
	}
	Broker struct {
		IP              string
		Port            int
		CertFile        string
		KeyFile         string
		AuthExchange    string
		AuthAllExchange string
	}
	Loggr struct {
		Push   bool
		Url    string
		ApiKey string
	}
	Librato struct {
		Push     bool
		Email    string
		Token    string
		Interval int
	}
	Opsview struct {
		Push bool
		Host string
	}
	Kontrold struct {
		Api struct {
			Port int
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
			SSLIPS  string
		}
		RabbitMq struct {
			Host     string
			Port     string
			Login    string
			Password string
			Vhost    string
		}
	}
	FollowFeed struct {
		Host          string
		Port          int
		ComponentUser string
		Password      string
		Vhost         string
	}
}

var FileProfile string
var PillarProfile string
var Profile string
var Current Config
var LogDebug bool
var Uuid string
var Host string
var BrokerDomain string
var Region string

func init() {
	flag.StringVar(&FileProfile, "c", "", "Configuration profile from file")
	flag.StringVar(&FileProfile, "config", "", "Alias for -c")
	flag.StringVar(&PillarProfile, "p", "", "Configuration profile from saltstack pillar")
	flag.BoolVar(&LogDebug, "d", false, "Log debug messages")
	flag.StringVar(&Uuid, "u", "", "Enable kontrol mode")
	flag.StringVar(&Host, "h", "", "Hostname to be resolved")
	flag.StringVar(&BrokerDomain, "a", "", "Send kontrol a custom domain istead of os.Hostname")
	flag.StringVar(&BrokerDomain, "domain", "", "Alias for -a")
	flag.StringVar(&Region, "r", "", "Region")

	flag.Parse()
	if flag.NArg() != 0 {
		flag.PrintDefaults()
		os.Exit(1)
	}
	if FileProfile == "" && PillarProfile == "" {
		fmt.Println("Please specify a configuration profile via -c or -p.")
		flag.PrintDefaults()
		os.Exit(1)
	}
	if FileProfile != "" && PillarProfile != "" {
		fmt.Println("The flags -c and -p are exclusive.")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var configCommand *exec.Cmd
	if FileProfile != "" {
		Profile = FileProfile
		configCommand = exec.Command("node", "-e", "require('koding-config-manager').printJson('main."+FileProfile+"')")
	}
	if PillarProfile != "" {
		Profile = PillarProfile
		configCommand = exec.Command("salt-call", "pillar.get", PillarProfile, "--output=json", "--log-level=warning")
	}

	configJSON, err := configCommand.CombinedOutput()
	if err != nil {
		fmt.Printf("Could not execute configuration source: %s\nConfiguration source output:\n%s\n", err.Error(), configJSON)
		os.Exit(1)
	}

	err = json.Unmarshal(configJSON, &Current)
	if err == nil && (Current == Config{}) {
		err = fmt.Errorf("Empty configuration.")
	}
	if err != nil {
		fmt.Printf("Could not unmarshal configuration: %s\nConfiguration source output:\n%s\n", err.Error(), configJSON)
		os.Exit(1)
	}

	sigChannel := make(chan os.Signal)
	signal.Notify(sigChannel, syscall.SIGUSR2)
	go func() {
		for _ = range sigChannel {
			LogDebug = !LogDebug
			fmt.Printf("config.LogDebug: %v\n", LogDebug)
		}
	}()

}
<|endoftext|>