<commit_before>package cas

import (
	"crypto/rand"
	"fmt"
	"net/http"
	"net/url"

	"github.com/golang/glog"
)

// Options : Client configuration options
type Options struct {
	URL          *url.URL     // URL to the CAS service
	Store        TicketStore  // Custom TicketStore, if nil a MemoryStore will be used
	Client       *http.Client // Custom http client to allow options for http connections
	SendService  bool         // Custom sendService to determine whether you need to send service param
	URLScheme    URLScheme    // Custom url scheme, can be used to modify the request urls for the client
	Cookie       *http.Cookie // http.Cookie options, uses Path, Domain, MaxAge, HttpOnly, & Secure
	SessionStore SessionStore
}

// Client implements the main protocol
type Client struct {
	tickets   TicketStore
	client    *http.Client
	urlScheme URLScheme
	cookie    *http.Cookie

	sessions    SessionStore
	sendService bool

	stValidator *ServiceTicketValidator
}

// NewClient creates a Client with the provided Options.
func NewClient(options *Options) *Client {
	if glog.V(2) {
		glog.Infof("cas: new client with options %v", options)
	}

	var tickets TicketStore
	if options.Store != nil {
		tickets = options.Store
	} else {
		tickets = &MemoryStore{}
	}

	var sessions SessionStore
	if options.SessionStore != nil {
		sessions = options.SessionStore
	} else {
		sessions = NewMemorySessionStore()
	}

	var urlScheme URLScheme
	if options.URLScheme != nil {
		urlScheme = options.URLScheme
	} else {
		urlScheme = NewDefaultURLScheme(options.URL)
	}

	var client *http.Client
	if options.Client != nil {
		client = options.Client
	} else {
		client = &http.Client{}
	}

	var cookie *http.Cookie
	if options.Cookie != nil {
		cookie = options.Cookie
	} else {
		cookie = &http.Cookie{
			MaxAge:   86400,
			HttpOnly: false,
			Secure:   false,
		}
	}

	return &Client{
		tickets:     tickets,
		client:      client,
		urlScheme:   urlScheme,
		cookie:      cookie,
		sessions:    sessions,
		sendService: options.SendService,
		stValidator: NewServiceTicketValidator(client, options.URL),
	}
}

// Handle wraps a http.Handler to provide CAS authentication for the handler.
func (c *Client) Handle(h http.Handler) http.Handler {
	return &clientHandler{
		c: c,
		h: h,
	}
}

// HandleFunc wraps a function to provide CAS authentication for the handler function.
func (c *Client) HandleFunc(h func(http.ResponseWriter, *http.Request)) http.Handler {
	return c.Handle(http.HandlerFunc(h))
}

// requestURL determines an absolute URL from the http.Request.
func requestURL(r *http.Request) (*url.URL, error) {
	u, err := url.Parse(r.URL.String())
	if err != nil {
		return nil, err
	}

	u.Host = r.Host
	if host := r.Header.Get("X-Forwarded-Host"); host != "" {
		u.Host = host
	}

	u.Scheme = "http"
	if scheme := r.Header.Get("X-Forwarded-Proto"); scheme != "" {
		u.Scheme = scheme
	} else if r.TLS != nil {
		u.Scheme = "https"
	}

	return u, nil
}

// LoginUrlForRequest determines the CAS login URL for the http.Request.
func (c *Client) LoginUrlForRequest(r *http.Request) (string, error) {
	u, err := c.urlScheme.Login()
	if err != nil {
		return "", err
	}

	service, err := requestURL(r)
	if err != nil {
		return "", err
	}

	q := u.Query()
	q.Add("service", sanitisedURLString(service))
	u.RawQuery = q.Encode()

	return u.String(), nil
}

// LogoutUrlForRequest determines the CAS logout URL for the http.Request.
func (c *Client) LogoutUrlForRequest(r *http.Request) (string, error) {
	u, err := c.urlScheme.Logout()
	if err != nil {
		return "", err
	}

	if c.sendService {
		service, err := requestURL(r)
		if err != nil {
			return "", err
		}

		q := u.Query()
		q.Add("service", sanitisedURLString(service))
		u.RawQuery = q.Encode()
	}

	return u.String(), nil
}

// ServiceValidateUrlForRequest determines the CAS serviceValidate URL for the ticket and http.Request.
func (c *Client) ServiceValidateUrlForRequest(ticket string, r *http.Request) (string, error) {
	service, err := requestURL(r)
	if err != nil {
		return "", err
	}
	return c.stValidator.ServiceValidateUrl(service, ticket)
}

// ValidateUrlForRequest determines the CAS validate URL for the ticket and http.Request.
func (c *Client) ValidateUrlForRequest(ticket string, r *http.Request) (string, error) {
	service, err := requestURL(r)
	if err != nil {
		return "", err
	}
	return c.stValidator.ValidateUrl(service, ticket)
}

// RedirectToLogout replies to the request with a redirect URL to log out of CAS.
func (c *Client) RedirectToLogout(w http.ResponseWriter, r *http.Request) {
	u, err := c.LogoutUrlForRequest(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if glog.V(2) {
		glog.Infof("Logging out, redirecting client to %v with status %v",
			u, http.StatusFound)
	}

	c.clearSession(w, r)
	http.Redirect(w, r, u, http.StatusFound)
}

// RedirectToLogin replies to the request with a redirect URL to authenticate with CAS.
func (c *Client) RedirectToLogin(w http.ResponseWriter, r *http.Request) {
	u, err := c.LoginUrlForRequest(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if glog.V(2) {
		glog.Infof("Redirecting client to %v with status %v", u, http.StatusFound)
	}

	http.Redirect(w, r, u, http.StatusFound)
}

// validateTicket performs CAS ticket validation with the given ticket and service.
func (c *Client) validateTicket(ticket string, service *http.Request) error {
	serviceURL, err := requestURL(service)
	if err != nil {
		return err
	}

	success, err := c.stValidator.ValidateTicket(serviceURL, ticket)
	if err != nil {
		return err
	}

	if err := c.tickets.Write(ticket, success); err != nil {
		return err
	}

	return nil
}

// getSession finds or creates a session for the request.
//
// A cookie is set on the response if one is not provided with the request.
// Validates the ticket if the URL parameter is provided.
func (c *Client) getSession(w http.ResponseWriter, r *http.Request) {
	cookie := c.getCookie(w, r)

	if s, ok := c.sessions.Get(cookie.Value); ok {
		if t, err := c.tickets.Read(s); err == nil {
			if glog.V(1) {
				glog.Infof("Re-used ticket %s for %s", s, t.User)
			}

			setAuthenticationResponse(r, t)
			return
		} else {
			if glog.V(2) {
				glog.Infof("Ticket %v not in %T: %v", s, c.tickets, err)
			}

			if glog.V(1) {
				glog.Infof("Clearing ticket %s, no longer exists in ticket store", s)
			}

			clearCookie(w, cookie)
		}
	}

	if ticket := r.URL.Query().Get("ticket"); ticket != "" {
		if err := c.validateTicket(ticket, r); err != nil {
			if glog.V(2) {
				glog.Infof("Error validating ticket: %v", err)
			}
			return // allow ServeHTTP()
		}

		c.setSession(cookie.Value, ticket)

		if t, err := c.tickets.Read(ticket); err == nil {
			if glog.V(1) {
				glog.Infof("Validated ticket %s for %s", ticket, t.User)
			}

			setAuthenticationResponse(r, t)
			return
		} else {
			if glog.V(2) {
				glog.Infof("Ticket %v not in %T: %v", ticket, c.tickets, err)
			}

			if glog.V(1) {
				glog.Infof("Clearing ticket %s, no longer exists in ticket store", ticket)
			}

			clearCookie(w, cookie)
		}
	}
}

// getCookie finds or creates the session cookie on the response.
func (c *Client) getCookie(w http.ResponseWriter, r *http.Request) *http.Cookie {
	cookie, err := r.Cookie(sessionCookieName)
	if err != nil {
		// NOTE: Intentionally not enabling HttpOnly so the cookie can
		//       still be used by Ajax requests.
		cookie = &http.Cookie{
			Name:     sessionCookieName,
			Value:    newSessionID(),
			Path:     c.cookie.Path,
			Domain:   c.cookie.Domain,
			MaxAge:   c.cookie.MaxAge,
			HttpOnly: c.cookie.HttpOnly,
			Secure:   c.cookie.Secure,
		}

		if glog.V(2) {
			glog.Infof("Setting %v cookie with value: %v", cookie.Name, cookie.Value)
		}

		r.AddCookie(cookie) // so we can find it later if required
		http.SetCookie(w, cookie)
	}

	return cookie
}

// newSessionId generates a new opaque session identifier for use in the cookie.
func newSessionID() string {
	const alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

	// generate 64 character string
	bytes := make([]byte, 64)
	rand.Read(bytes)

	for k, v := range bytes {
		bytes[k] = alphabet[v%byte(len(alphabet))]
	}

	return string(bytes)
}

// clearCookie invalidates and removes the cookie from the client.
func clearCookie(w http.ResponseWriter, c *http.Cookie) {
	c.MaxAge = -1
	http.SetCookie(w, c)
}

// setSession stores the session id to ticket mapping in the Client.
func (c *Client) setSession(id string, ticket string) {
	if glog.V(2) {
		glog.Infof("Recording session, %v -> %v", id, ticket)
	}

	c.sessions.Set(id, ticket)
}

// clearSession removes the session from the client and clears the cookie.
func (c *Client) clearSession(w http.ResponseWriter, r *http.Request) {
	cookie := c.getCookie(w, r)

	if s, ok := c.sessions.Get(cookie.Value); ok {
		if err := c.tickets.Delete(s); err != nil {
			fmt.Printf("Failed to remove %v from %T: %v\n", cookie.Value, c.tickets, err)
			if glog.V(2) {
				glog.Errorf("Failed to remove %v from %T: %v", cookie.Value, c.tickets, err)
			}
		}

		c.deleteSession(s)
	}

	clearCookie(w, cookie)
}

// deleteSession removes the session from the client
func (c *Client) deleteSession(id string) {
	c.sessions.Delete(id)
}
<commit_msg>fix clear session passed wrong argument<commit_after>package cas

import (
	"crypto/rand"
	"fmt"
	"net/http"
	"net/url"

	"github.com/golang/glog"
)

// Options : Client configuration options
type Options struct {
	URL          *url.URL     // URL to the CAS service
	Store        TicketStore  // Custom TicketStore, if nil a MemoryStore will be used
	Client       *http.Client // Custom http client to allow options for http connections
	SendService  bool         // Custom sendService to determine whether you need to send service param
	URLScheme    URLScheme    // Custom url scheme, can be used to modify the request urls for the client
	Cookie       *http.Cookie // http.Cookie options, uses Path, Domain, MaxAge, HttpOnly, & Secure
	SessionStore SessionStore
}

// Client implements the main protocol
type Client struct {
	tickets   TicketStore
	client    *http.Client
	urlScheme URLScheme
	cookie    *http.Cookie

	sessions    SessionStore
	sendService bool

	stValidator *ServiceTicketValidator
}

// NewClient creates a Client with the provided Options.
func NewClient(options *Options) *Client {
	if glog.V(2) {
		glog.Infof("cas: new client with options %v", options)
	}

	var tickets TicketStore
	if options.Store != nil {
		tickets = options.Store
	} else {
		tickets = &MemoryStore{}
	}

	var sessions SessionStore
	if options.SessionStore != nil {
		sessions = options.SessionStore
	} else {
		sessions = NewMemorySessionStore()
	}

	var urlScheme URLScheme
	if options.URLScheme != nil {
		urlScheme = options.URLScheme
	} else {
		urlScheme = NewDefaultURLScheme(options.URL)
	}

	var client *http.Client
	if options.Client != nil {
		client = options.Client
	} else {
		client = &http.Client{}
	}

	var cookie *http.Cookie
	if options.Cookie != nil {
		cookie = options.Cookie
	} else {
		cookie = &http.Cookie{
			MaxAge:   86400,
			HttpOnly: false,
			Secure:   false,
		}
	}

	return &Client{
		tickets:     tickets,
		client:      client,
		urlScheme:   urlScheme,
		cookie:      cookie,
		sessions:    sessions,
		sendService: options.SendService,
		stValidator: NewServiceTicketValidator(client, options.URL),
	}
}

// Handle wraps a http.Handler to provide CAS authentication for the handler.
func (c *Client) Handle(h http.Handler) http.Handler {
	return &clientHandler{
		c: c,
		h: h,
	}
}

// HandleFunc wraps a function to provide CAS authentication for the handler function.
func (c *Client) HandleFunc(h func(http.ResponseWriter, *http.Request)) http.Handler {
	return c.Handle(http.HandlerFunc(h))
}

// requestURL determines an absolute URL from the http.Request.
func requestURL(r *http.Request) (*url.URL, error) {
	u, err := url.Parse(r.URL.String())
	if err != nil {
		return nil, err
	}

	u.Host = r.Host
	if host := r.Header.Get("X-Forwarded-Host"); host != "" {
		u.Host = host
	}

	u.Scheme = "http"
	if scheme := r.Header.Get("X-Forwarded-Proto"); scheme != "" {
		u.Scheme = scheme
	} else if r.TLS != nil {
		u.Scheme = "https"
	}

	return u, nil
}

// LoginUrlForRequest determines the CAS login URL for the http.Request.
func (c *Client) LoginUrlForRequest(r *http.Request) (string, error) {
	u, err := c.urlScheme.Login()
	if err != nil {
		return "", err
	}

	service, err := requestURL(r)
	if err != nil {
		return "", err
	}

	q := u.Query()
	q.Add("service", sanitisedURLString(service))
	u.RawQuery = q.Encode()

	return u.String(), nil
}

// LogoutUrlForRequest determines the CAS logout URL for the http.Request.
func (c *Client) LogoutUrlForRequest(r *http.Request) (string, error) {
	u, err := c.urlScheme.Logout()
	if err != nil {
		return "", err
	}

	if c.sendService {
		service, err := requestURL(r)
		if err != nil {
			return "", err
		}

		q := u.Query()
		q.Add("service", sanitisedURLString(service))
		u.RawQuery = q.Encode()
	}

	return u.String(), nil
}

// ServiceValidateUrlForRequest determines the CAS serviceValidate URL for the ticket and http.Request.
func (c *Client) ServiceValidateUrlForRequest(ticket string, r *http.Request) (string, error) {
	service, err := requestURL(r)
	if err != nil {
		return "", err
	}
	return c.stValidator.ServiceValidateUrl(service, ticket)
}

// ValidateUrlForRequest determines the CAS validate URL for the ticket and http.Request.
func (c *Client) ValidateUrlForRequest(ticket string, r *http.Request) (string, error) {
	service, err := requestURL(r)
	if err != nil {
		return "", err
	}
	return c.stValidator.ValidateUrl(service, ticket)
}

// RedirectToLogout replies to the request with a redirect URL to log out of CAS.
func (c *Client) RedirectToLogout(w http.ResponseWriter, r *http.Request) {
	u, err := c.LogoutUrlForRequest(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if glog.V(2) {
		glog.Infof("Logging out, redirecting client to %v with status %v",
			u, http.StatusFound)
	}

	c.clearSession(w, r)
	http.Redirect(w, r, u, http.StatusFound)
}

// RedirectToLogin replies to the request with a redirect URL to authenticate with CAS.
func (c *Client) RedirectToLogin(w http.ResponseWriter, r *http.Request) {
	u, err := c.LoginUrlForRequest(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if glog.V(2) {
		glog.Infof("Redirecting client to %v with status %v", u, http.StatusFound)
	}

	http.Redirect(w, r, u, http.StatusFound)
}

// validateTicket performs CAS ticket validation with the given ticket and service.
func (c *Client) validateTicket(ticket string, service *http.Request) error {
	serviceURL, err := requestURL(service)
	if err != nil {
		return err
	}

	success, err := c.stValidator.ValidateTicket(serviceURL, ticket)
	if err != nil {
		return err
	}

	if err := c.tickets.Write(ticket, success); err != nil {
		return err
	}

	return nil
}

// getSession finds or creates a session for the request.
//
// A cookie is set on the response if one is not provided with the request.
// Validates the ticket if the URL parameter is provided.
func (c *Client) getSession(w http.ResponseWriter, r *http.Request) {
	cookie := c.getCookie(w, r)

	if s, ok := c.sessions.Get(cookie.Value); ok {
		if t, err := c.tickets.Read(s); err == nil {
			if glog.V(1) {
				glog.Infof("Re-used ticket %s for %s", s, t.User)
			}

			setAuthenticationResponse(r, t)
			return
		} else {
			if glog.V(2) {
				glog.Infof("Ticket %v not in %T: %v", s, c.tickets, err)
			}

			if glog.V(1) {
				glog.Infof("Clearing ticket %s, no longer exists in ticket store", s)
			}

			clearCookie(w, cookie)
		}
	}

	if ticket := r.URL.Query().Get("ticket"); ticket != "" {
		if err := c.validateTicket(ticket, r); err != nil {
			if glog.V(2) {
				glog.Infof("Error validating ticket: %v", err)
			}
			return // allow ServeHTTP()
		}

		c.setSession(cookie.Value, ticket)

		if t, err := c.tickets.Read(ticket); err == nil {
			if glog.V(1) {
				glog.Infof("Validated ticket %s for %s", ticket, t.User)
			}

			setAuthenticationResponse(r, t)
			return
		} else {
			if glog.V(2) {
				glog.Infof("Ticket %v not in %T: %v", ticket, c.tickets, err)
			}

			if glog.V(1) {
				glog.Infof("Clearing ticket %s, no longer exists in ticket store", ticket)
			}

			clearCookie(w, cookie)
		}
	}
}

// getCookie finds or creates the session cookie on the response.
func (c *Client) getCookie(w http.ResponseWriter, r *http.Request) *http.Cookie {
	cookie, err := r.Cookie(sessionCookieName)
	if err != nil {
		// NOTE: Intentionally not enabling HttpOnly so the cookie can
		//       still be used by Ajax requests.
		cookie = &http.Cookie{
			Name:     sessionCookieName,
			Value:    newSessionID(),
			Path:     c.cookie.Path,
			Domain:   c.cookie.Domain,
			MaxAge:   c.cookie.MaxAge,
			HttpOnly: c.cookie.HttpOnly,
			Secure:   c.cookie.Secure,
		}

		if glog.V(2) {
			glog.Infof("Setting %v cookie with value: %v", cookie.Name, cookie.Value)
		}

		r.AddCookie(cookie) // so we can find it later if required
		http.SetCookie(w, cookie)
	}

	return cookie
}

// newSessionId generates a new opaque session identifier for use in the cookie.
func newSessionID() string {
	const alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

	// generate 64 character string
	bytes := make([]byte, 64)
	rand.Read(bytes)

	for k, v := range bytes {
		bytes[k] = alphabet[v%byte(len(alphabet))]
	}

	return string(bytes)
}

// clearCookie invalidates and removes the cookie from the client.
func clearCookie(w http.ResponseWriter, c *http.Cookie) {
	c.MaxAge = -1
	http.SetCookie(w, c)
}

// setSession stores the session id to ticket mapping in the Client.
func (c *Client) setSession(id string, ticket string) {
	if glog.V(2) {
		glog.Infof("Recording session, %v -> %v", id, ticket)
	}

	c.sessions.Set(id, ticket)
}

// clearSession removes the session from the client and clears the cookie.
func (c *Client) clearSession(w http.ResponseWriter, r *http.Request) {
	cookie := c.getCookie(w, r)

	if serviceTicket, ok := c.sessions.Get(cookie.Value); ok {
		if err := c.tickets.Delete(serviceTicket); err != nil {
			fmt.Printf("Failed to remove %v from %T: %v\n", cookie.Value, c.tickets, err)
			if glog.V(2) {
				glog.Errorf("Failed to remove %v from %T: %v", cookie.Value, c.tickets, err)
			}
		}

		c.deleteSession(cookie.Value)
	}

	clearCookie(w, cookie)
}

// deleteSession removes the session from the client
func (c *Client) deleteSession(id string) {
	c.sessions.Delete(id)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"
	"flag"

	"code.google.com/p/go.net/websocket"
)

const BURST_SIZE = 50
const BURST_INTV = 500
const CONNECTIONS = 20000
const SERVICE = 8080
const HOST = "127.0.0.1"
const ORIGIN = "127.0.0.1"

var p = Pool {
	connections: make(map[*Connection]bool),
	subscribe: make(chan *Connection),
	unsubscribe: make(chan *Connection),
}

func listen(conn *Connection) {
	p.subscribe <- conn
	for {
		var message string
		err := websocket.Message.Receive(conn.ws, &message)
		if err != nil {
			fmt.Printf("Receive: %s\n", err)
			break
		}
	}
	p.unsubscribe <- conn
	conn.ws.Close()
}

func client(id int, host string, service int, origin string) {
	orig := fmt.Sprintf("http://%s/", origin)
	targ := fmt.Sprintf("ws://%s:%d/", host, service)

	ws, err := websocket.Dial(targ, "", orig)
	if err != nil {
		fmt.Printf("Dial: %s\n", err)
		return
	}
	var conn = &Connection {
		ws: ws,
        }
	go listen(conn)
}

func flood(
	connid *int,
	connections int,
	host string,
	service int,
	origin string,
	burst_size int,
	burst_intv int) {
	for {
		for y := 0; y < burst_size; y++ {
			if *connid >= connections {
				break
			}
			client(*connid, host, service, origin)
			*connid++
		}
		time.Sleep(time.Duration(burst_intv) * time.Millisecond)
	}
}

func main() {
	var connections = flag.Int("connections", CONNECTIONS, "Number of concurent connections")
	var host = flag.String("host", HOST, "Server host")
	var port = flag.Int("port", SERVICE, "Server port")
	var origin = flag.String("origin", HOST, "Client origin")
	var burst_size = flag.Int("burst-size", BURST_SIZE, "Number of concurent connections per bucket send")
	var burst_intv = flag.Int("burst-intv", BURST_INTV, "Interval between each bust of connections")

	flag.Parse()

	var connid = 0
	go p.Dispatch()
	go flood(&connid, *connections, *host, *port, *origin, *burst_size, *burst_intv);

	fmt.Printf("Wait for initialization...\n")
	for {
		if len(p.connections) >= 1 {
			break
		}
		time.Sleep(200 * time.Millisecond)
	}

	for {
		time.Sleep(1 * time.Second)
		fmt.Printf("Connections spawned: %d\n", connid)
		if len(p.connections) <= 0 {
			break
		}
	}
}
<commit_msg>Review default options<commit_after>package main

import (
	"fmt"
	"time"
	"flag"

	"code.google.com/p/go.net/websocket"
)

const BURST_SIZE = 10
const BURST_INTV = 100
const CONNECTIONS = 1000
const SERVICE = 8080
const HOST = "127.0.0.1"
const ORIGIN = "127.0.0.1"

var p = Pool {
	connections: make(map[*Connection]bool),
	subscribe: make(chan *Connection),
	unsubscribe: make(chan *Connection),
}

func listen(conn *Connection) {
	p.subscribe <- conn
	for {
		var message string
		err := websocket.Message.Receive(conn.ws, &message)
		if err != nil {
			fmt.Printf("Receive: %s\n", err)
			break
		}
	}
	p.unsubscribe <- conn
	conn.ws.Close()
}

func client(id int, host string, service int, origin string) {
	orig := fmt.Sprintf("http://%s/", origin)
	targ := fmt.Sprintf("ws://%s:%d/", host, service)

	ws, err := websocket.Dial(targ, "", orig)
	if err != nil {
		fmt.Printf("Dial: %s\n", err)
		return
	}
	var conn = &Connection {
		ws: ws,
        }
	go listen(conn)
}

func flood(
	connid *int,
	connections int,
	host string,
	service int,
	origin string,
	burst_size int,
	burst_intv int) {
	for {
		for y := 0; y < burst_size; y++ {
			if *connid >= connections {
				break
			}
			client(*connid, host, service, origin)
			*connid++
		}
		time.Sleep(time.Duration(burst_intv) * time.Millisecond)
	}
}

func main() {
	var connections = flag.Int("connections", CONNECTIONS, "Number of concurent connections")
	var host = flag.String("host", HOST, "Server host")
	var port = flag.Int("port", SERVICE, "Server port")
	var origin = flag.String("origin", HOST, "Client origin")
	var burst_size = flag.Int("burst-size", BURST_SIZE, "Number of concurent connections per bucket send")
	var burst_intv = flag.Int("burst-intv", BURST_INTV, "Interval between each bust of connections")

	flag.Parse()

	var connid = 0
	go p.Dispatch()
	go flood(&connid, *connections, *host, *port, *origin, *burst_size, *burst_intv);

	fmt.Printf("Wait for initialization...\n")
	for {
		if len(p.connections) >= 1 {
			break
		}
		time.Sleep(200 * time.Millisecond)
	}

	for {
		time.Sleep(1 * time.Second)
		fmt.Printf("Connections spawned: %d\n", connid)
		if len(p.connections) <= 0 {
			break
		}
	}
}
<|endoftext|>


<commit_before>package grestclient

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	rt "reflect"
)

type Client interface {

	//Headers returns the default headers that will
	//be set with every request made with the client.
	Headers() http.Header

	//SetHeaders sets the default headers that will be sent with
	//every request made with this client.
	SetHeaders(http.Header)

	//Query returns the default query to use for all requests
	Query() url.Values

	//SetQuery sets the default query to use for all requests
	SetQuery(url.Values)

	//SetBaseUrl sets the base url to use for all requests
	//If you want to use a different temporarily it is best to
	//create a new client with the new base url. Call
	//Clone() to get a clone of this
	//client's settings and then change the url on the clone.
	//An error should be returned if the url is "unsupported" by
	//the implementation. For example, "unix://tmp.soc".
	//Any query parameters added here should be ignored.
	//Clients should use the SetQuery method to set default
	//query parameters
	SetBaseUrl(*url.URL) error

	//BaseUrl returns the base url being used. This implementation
	//allows you to change the base url here directly but other
	//implementations might give you a clone so changing it won't affect
	//the client. In those cases, use SetBaseUrl to change the url.
	BaseUrl() *url.URL

	//Clones the client with everything the old client had
	//Ideally, clones should be independent of the original and can be changed
	//without affecting the original and vice versa.
	//This implementation, however, shared the http.Client among clones.
	//All other 'things' like headers, base url, query, marshalers are separate and
	//can be adjusted without affecting the original/clones.
	Clone() Client

	//GetHttpClient returns the current http.Client being used
	//If none has been set, this should return http.DefaultClient
	GetHttpClient() *http.Client

	//SetHttpClient sets the http.Client to use during requests
	//Use this to customize your http.Client as you wish. If you
	//don't set one, the default http.Client will be used.
	SetHttpClient(*http.Client)

	//SetMarshaler sets the marshal function to be used
	//to marshal the request bodies for requests
	//Doesn't have to mirror the Unmarshaler. Send plain text, get back json
	//Default is a json marshaler
	SetMarshaler(MarshalerFunc)
	//SetUnmarshaler sets the unmarshal function to be used
	//to unmarshal the response body for responses
	//Doesn't have to mirror the Marshaler. Send XML, get back json
	//Default is a json unmarshaler
	SetUnmarshaler(UnmarshalerFunc)

	//AddRequestMutator adds a mutator that the request will be
	//passed through before executing the request. All RequestMutators
	//are called AFTER the Marshaler is used.
	//RequestMutators should be called in the order they were added
	AddRequestMutators(...RequestMutator) Client

	//AddResponseMutator adds a mutator that the response will be
	//passed through after the server responds. All ResponseMutators
	//are called BEFORE the Unmarshaler is used.
	//ResponseMutators should be called in the order they were added
	AddResponseMutators(...ResponseMutator) Client

	//RemoveRequestMutator removes a request mutator
	SetRequestMutators(...RequestMutator) Client
	//RemoveResponseMutator removes a response mutator
	SetResponseMutators(...ResponseMutator) Client

	//Returns the RequestMutators
	RequestMutators() []RequestMutator
	//Returns the ResponseMutators
	ResponseMutators() []ResponseMutator

	//Get performs a get request with the base url plus the path appended to it.
	//You can send query values, header values and
	//supply a successResult that will be populated if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Get(path string, headers http.Header, query url.Values, successResult interface{}, errorResult interface{}) (*http.Response, error)

	//Post performs a post request with the base url plus the path appended to it.
	//You can send query values, header values and
	//supply a successResult that will be populated if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//With post you can also provide a post body.
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Post(path string, headers http.Header, query url.Values, postBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error)

	//Put performs a put request with the base url plus the path appended to it.
	//You can send query values, header values and
	//supply a successResult that will be populated if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//With put you can also provide a put body.
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Put(path string, headers http.Header, query url.Values, putBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error)

	//Patch performs a patch request with the base url plus the path appended to it.
	//You can send query values, header values and
	//supply a successResult that will be populated if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//With patch you can also provide a patch body.
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Patch(path string, headers http.Header, query url.Values, patchBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error)

	//Head performs a head request with the base url plus the path appended to it.
	//You can send header values and supply a successResult that will be populated
	//if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Head(path string, headers http.Header, successResult interface{}, errorResultg interface{}) (*http.Response, error)

	//Option performs an option request with the base url plus the path appended to it.
	//You can send header values and supply a successResult that will be populated
	//if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Options(path string, headers http.Header, successResult interface{}, errorResult interface{}) (*http.Response, error)

	//Delete performs an delete request with the base url plus the path appended to it.
	//You can send header values and supply a successResult that will be populated
	//if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Delete(path string, headers http.Header, query url.Values, successResult interface{}, errorResult interface{}) (*http.Response, error)
}

type ReadLener interface {
	io.Reader
	Len() int
}

//MarshalerFunc takes something and converts it into a
//ReadLener that can be used for the request body
type MarshalerFunc func(v interface{}) (ReadLener, error)

//UnmarshalerFunc takes the response body and converts it into
//something you can use.
type UnmarshalerFunc func(b io.ReadCloser, v interface{}) error

//ByteSliceToReadCloser takes a byte slice and converts it to an
//ReadLener that can be used as a request/resonse body
func ByteSliceToReadLener(b []byte) (ReadLener, error) {
	if b == nil {
		return nil, errors.New("ReadCloserFromByteSlice received a nil byte slice.")
	}

	buf := bytes.NewBuffer(b)
	return buf, nil
}

//StringToReadCloser takes a string and converts it to an
//ReadLener that can be used as a request/resonse body
func StringToReadLener(s string) ReadLener {
	buf := bytes.NewBufferString(s)
	return buf
}

//JsonMarshalerFunc can be used to marshal request bodies
//into json
func JsonMarshalerFunc(v interface{}) (ReadLener, error) {
	b, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	body, err := ByteSliceToReadLener(b)
	if err != nil {
		return nil, err
	}
	return body, nil
}

//JsonUnmarshalerFunc can be used to unmarshal response bodies
//from json
func JsonUnmarshalerFunc(body io.ReadCloser, v interface{}) error {

	b, err := ioutil.ReadAll(body)
	if err != nil {
		return err
	}

	err = json.Unmarshal(b, v)

	if err != nil {
		return err
	}

	return nil
}

//StringMarshalerFunc can be used to marshal strings into a request.
func StringMarshalerFunc(v interface{}) (ReadLener, error) {
	switch t := v.(type) {
	case fmt.Stringer:
		return StringToReadLener(t.String()), nil
	case string:
		return StringToReadLener(t), nil
	}

	return nil, errors.New("Did not know how to use the body as text.")
}

//StringUnmarshalerFunc can be used to unmarshal strings from a response.
func StringUnmarshalerFunc(body io.ReadCloser, v interface{}) error {
	if v == nil {
		return nil
	}

	b, err := ioutil.ReadAll(body)

	if err != nil {
		return err
	}

	switch v.(type) {
	case *string:
		elem := rt.ValueOf(v).Elem()

		if elem.CanSet() {
			elem.SetString(string(b))
			return nil
		}
	case string:
		return errors.New("You must pass the string by reference: ")
	}

	return errors.New("Did not know how to unmarshal the text coming back.")
}

//JsonContentTypeMutator sets the Content-Type of the request to be
//application/json
func JsonContentTypeMutator(r *http.Request) error {
	r.Header.Add("Content-Type", "application/json")
	return nil
}

func JsonAcceptMutator(r *http.Request) error {
	r.Header.Add("Accept", "application/json")
	return nil
}

//RequestMutators are called before the request is made but after the marshaler function has been
//called.
type RequestMutator func(*http.Request) error
type ResponseMutator func(*http.Response) error

//SetupClientForJson is a convenience method that sets the
//marshaler and unmarshaler funcs on the client to be the
//Json funcs in this package. It also sets a request mutator
//to set the Content-Type to json.
func SetupForJson(c Client) {
	c.SetMarshaler(JsonMarshalerFunc)
	c.SetUnmarshaler(JsonUnmarshalerFunc)
	c.AddRequestMutators(JsonContentTypeMutator)
	c.AddRequestMutators(JsonAcceptMutator)
}
<commit_msg>changed spacing.<commit_after>package grestclient

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	rt "reflect"
)

type Client interface {

	//Headers returns the default headers that will
	//be set with every request made with the client.
	Headers() http.Header

	//SetHeaders sets the default headers that will be sent with
	//every request made with this client.
	SetHeaders(http.Header)

	//Query returns the default query to use for all requests
	Query() url.Values

	//SetQuery sets the default query to use for all requests
	SetQuery(url.Values)

	//SetBaseUrl sets the base url to use for all requests
	//If you want to use a different temporarily it is best to
	//create a new client with the new base url. Call
	//Clone() to get a clone of this
	//client's settings and then change the url on the clone.
	//An error should be returned if the url is "unsupported" by
	//the implementation. For example, "unix://tmp.soc".
	//Any query parameters added here should be ignored.
	//Clients should use the SetQuery method to set default
	//query parameters
	SetBaseUrl(*url.URL) error

	//BaseUrl returns the base url being used. This implementation
	//allows you to change the base url here directly but other
	//implementations might give you a clone so changing it won't affect
	//the client. In those cases, use SetBaseUrl to change the url.
	BaseUrl() *url.URL

	//Clones the client with everything the old client had
	//Ideally, clones should be independent of the original and can be changed
	//without affecting the original and vice versa.
	//This implementation, however, shared the http.Client among clones.
	//All other 'things' like headers, base url, query, marshalers are separate and
	//can be adjusted without affecting the original/clones.
	Clone() Client

	//GetHttpClient returns the current http.Client being used
	//If none has been set, this should return http.DefaultClient
	GetHttpClient() *http.Client

	//SetHttpClient sets the http.Client to use during requests
	//Use this to customize your http.Client as you wish. If you
	//don't set one, the default http.Client will be used.
	SetHttpClient(*http.Client)

	//SetMarshaler sets the marshal function to be used
	//to marshal the request bodies for requests
	//Doesn't have to mirror the Unmarshaler. Send plain text, get back json
	//Default is a json marshaler
	SetMarshaler(MarshalerFunc)
	//SetUnmarshaler sets the unmarshal function to be used
	//to unmarshal the response body for responses
	//Doesn't have to mirror the Marshaler. Send XML, get back json
	//Default is a json unmarshaler
	SetUnmarshaler(UnmarshalerFunc)

	//AddRequestMutator adds a mutator that the request will be
	//passed through before executing the request. All RequestMutators
	//are called AFTER the Marshaler is used.
	//RequestMutators should be called in the order they were added
	AddRequestMutators(...RequestMutator) Client

	//AddResponseMutator adds a mutator that the response will be
	//passed through after the server responds. All ResponseMutators
	//are called BEFORE the Unmarshaler is used.
	//ResponseMutators should be called in the order they were added
	AddResponseMutators(...ResponseMutator) Client

	//RemoveRequestMutator removes a request mutator
	SetRequestMutators(...RequestMutator) Client
	//RemoveResponseMutator removes a response mutator
	SetResponseMutators(...ResponseMutator) Client

	//Returns the RequestMutators
	RequestMutators() []RequestMutator
	//Returns the ResponseMutators
	ResponseMutators() []ResponseMutator

	//Get performs a get request with the base url plus the path appended to it.
	//You can send query values, header values and
	//supply a successResult that will be populated if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Get(
		path string,
		headers http.Header,
		query url.Values,
		successResult interface{},
		errorResult interface{},
	) (*http.Response, error)

	//Post performs a post request with the base url plus the path appended to it.
	//You can send query values, header values and
	//supply a successResult that will be populated if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//With post you can also provide a post body.
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Post(
		path string,
		headers http.Header,
		query url.Values,
		postBody interface{},
		successResult interface{},
		errorResult interface{},
	) (*http.Response, error)

	//Put performs a put request with the base url plus the path appended to it.
	//You can send query values, header values and
	//supply a successResult that will be populated if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//With put you can also provide a put body.
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Put(
		path string,
		headers http.Header,
		query url.Values,
		putBody interface{},
		successResult interface{},
		errorResult interface{},
	) (*http.Response, error)

	//Patch performs a patch request with the base url plus the path appended to it.
	//You can send query values, header values and
	//supply a successResult that will be populated if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//With patch you can also provide a patch body.
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Patch(
		path string,
		headers http.Header,
		query url.Values,
		patchBody interface{},
		successResult interface{},
		errorResult interface{},
	) (*http.Response, error)

	//Head performs a head request with the base url plus the path appended to it.
	//You can send header values and supply a successResult that will be populated
	//if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Head(path string, headers http.Header, successResult interface{}, errorResultg interface{}) (*http.Response, error)

	//Option performs an option request with the base url plus the path appended to it.
	//You can send header values and supply a successResult that will be populated
	//if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Options(
		path string,
		headers http.Header,
		successResult interface{},
		errorResult interface{},
	) (*http.Response, error)

	//Delete performs an delete request with the base url plus the path appended to it.
	//You can send header values and supply a successResult that will be populated
	//if the http response has a return code less than 400.
	//errorResult is populated if the error code is 400 or more
	//Returns the raw http.Response and error similar to Do method of http.Client
	//The returned http.Response might be non-nil even though an error was also returned
	//depending on where the operation failed.
	Delete(
		path string,
		headers http.Header,
		query url.Values,
		successResult interface{},
		errorResult interface{},
	) (*http.Response, error)
}

type ReadLener interface {
	io.Reader
	Len() int
}

//MarshalerFunc takes something and converts it into a
//ReadLener that can be used for the request body
type MarshalerFunc func(v interface{}) (ReadLener, error)

//UnmarshalerFunc takes the response body and converts it into
//something you can use.
type UnmarshalerFunc func(b io.ReadCloser, v interface{}) error

//ByteSliceToReadCloser takes a byte slice and converts it to an
//ReadLener that can be used as a request/resonse body
func ByteSliceToReadLener(b []byte) (ReadLener, error) {
	if b == nil {
		return nil, errors.New("ReadCloserFromByteSlice received a nil byte slice.")
	}

	buf := bytes.NewBuffer(b)
	return buf, nil
}

//StringToReadCloser takes a string and converts it to an
//ReadLener that can be used as a request/resonse body
func StringToReadLener(s string) ReadLener {
	buf := bytes.NewBufferString(s)
	return buf
}

//JsonMarshalerFunc can be used to marshal request bodies
//into json
func JsonMarshalerFunc(v interface{}) (ReadLener, error) {
	b, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	body, err := ByteSliceToReadLener(b)
	if err != nil {
		return nil, err
	}
	return body, nil
}

//JsonUnmarshalerFunc can be used to unmarshal response bodies
//from json
func JsonUnmarshalerFunc(body io.ReadCloser, v interface{}) error {

	b, err := ioutil.ReadAll(body)
	if err != nil {
		return err
	}

	err = json.Unmarshal(b, v)

	if err != nil {
		return err
	}

	return nil
}

//StringMarshalerFunc can be used to marshal strings into a request.
func StringMarshalerFunc(v interface{}) (ReadLener, error) {
	switch t := v.(type) {
	case fmt.Stringer:
		return StringToReadLener(t.String()), nil
	case string:
		return StringToReadLener(t), nil
	}

	return nil, errors.New("Did not know how to use the body as text.")
}

//StringUnmarshalerFunc can be used to unmarshal strings from a response.
func StringUnmarshalerFunc(body io.ReadCloser, v interface{}) error {
	if v == nil {
		return nil
	}

	b, err := ioutil.ReadAll(body)

	if err != nil {
		return err
	}

	switch v.(type) {
	case *string:
		elem := rt.ValueOf(v).Elem()

		if elem.CanSet() {
			elem.SetString(string(b))
			return nil
		}
	case string:
		return errors.New("You must pass the string by reference: ")
	}

	return errors.New("Did not know how to unmarshal the text coming back.")
}

//JsonContentTypeMutator sets the Content-Type of the request to be
//application/json
func JsonContentTypeMutator(r *http.Request) error {
	r.Header.Add("Content-Type", "application/json")
	return nil
}

func JsonAcceptMutator(r *http.Request) error {
	r.Header.Add("Accept", "application/json")
	return nil
}

//RequestMutators are called before the request is made but after the marshaler function has been
//called.
type RequestMutator func(*http.Request) error
type ResponseMutator func(*http.Response) error

//SetupClientForJson is a convenience method that sets the
//marshaler and unmarshaler funcs on the client to be the
//Json funcs in this package. It also sets a request mutator
//to set the Content-Type to json.
func SetupForJson(c Client) {
	c.SetMarshaler(JsonMarshalerFunc)
	c.SetUnmarshaler(JsonUnmarshalerFunc)
	c.AddRequestMutators(JsonContentTypeMutator)
	c.AddRequestMutators(JsonAcceptMutator)
}
<|endoftext|>


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

package rollbar

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"time"

	"github.com/pkg/errors"
	api "github.com/zchee/go-rollbar/api/v1"
	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

type client struct {
	debugClient    *httpClient
	infoClient     *httpClient
	errorClient    *httpClient
	warnClient     *httpClient
	criticalClient *httpClient
}

type httpClient struct {
	token        string
	client       *http.Client
	endpoint     string
	debug        bool
	logger       Logger
	environment  string
	platform     string
	codeVersion  string
	serverHost   string
	serverRoot   string
	serverBranch string
}

var defaultHTTPClient = httpClient{
	client:      http.DefaultClient,
	endpoint:    api.DefaultEndpoint,
	logger:      nilLogger{},
	environment: "development",
	platform:    runtime.GOOS,
}

// New creates a new REST rollbar API client.
//
// The `token` is required, other optional parameters can be passed using the
// various `With...` functions.
func New(token string, options ...Option) Client {
	cl := defaultHTTPClient
	cl.token = token
	if debug, err := strconv.ParseBool(os.Getenv("ROLLBAR_DEBUG")); err == nil && debug {
		cl.debug = debug
	}

	for _, o := range options {
		o(&cl)
	}
	if _, ok := cl.logger.(nilLogger); cl.debug && ok {
		cl.logger = traceLogger{os.Stderr}
	}
	if cl.serverHost == "" {
		cl.serverHost, _ = os.Hostname()
	}

	return &client{
		debugClient:    &cl,
		infoClient:     &cl,
		errorClient:    &cl,
		warnClient:     &cl,
		criticalClient: &cl,
	}
}

// payload creates the rollbar payload data.
func (c *httpClient) payload(level Level, err error) *api.Payload {
	title := "<nil>"
	if err != nil {
		title = err.Error()
	}
	stack := CreateStack(3)

	data := &api.Data{
		Environment: c.environment,
		Body:        errorBody(err, stack),
		Level:       string(level),
		Timestamp:   time.Now().Unix(),
		Platform:    c.platform,
		Language:    language,
		Server: &api.Server{
			Host:   c.serverHost,
			Root:   c.serverRoot,
			Branch: c.serverBranch,
		},
		Fingerprint: stack.Fingerprint(),
		Title:       title,
		Notifier: &api.Notifier{
			Name:    Name,
			Version: Version,
		},
	}

	return &api.Payload{
		AccessToken: c.token,
		Data:        data,
	}
}

// post posts payload to rollbar.
func (c *httpClient) post(pctx context.Context, payload *api.Payload) (*api.Response, error) {
	if c.token == "" {
		return nil, errors.New("empty token")
	}

	data, err := json.Marshal(payload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to encode payload")
	}
	if c.debug {
		out, _ := json.MarshalIndent(payload, "", "  ")
		c.logger.Debugf(pctx, string(out)+"\n")
	}

	req, err := http.NewRequest(http.MethodPost, c.endpoint, bytes.NewReader(data))
	if err != nil {
		return nil, errors.Wrap(err, "failed to create new POST request")
	}

	ctx, cancel := context.WithCancel(pctx)
	defer cancel()

	req.Header.Set("Content-Type", "application/json")
	resp, err := ctxhttp.Do(ctx, c.client, req)
	if err != nil {
		return nil, errors.Wrap(err, "failed to POST to rollbar")
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, errors.Errorf("received response: %s", resp.Status)
	}

	return c.parseResponse(ctx, resp.Body), nil
}

// parseResponse parses the rollbar API response.
func (c *httpClient) parseResponse(ctx context.Context, rdr io.Reader) *api.Response {
	buf := new(bytes.Buffer)
	io.Copy(buf, rdr)

	c.logger.Debugf(ctx, "-----> %s (response)\n", c.endpoint)
	m := new(api.Response)
	if err := json.Unmarshal(buf.Bytes(), m); err != nil {
		c.logger.Debugf(ctx, "failed to unmarshal payload: %s\n", err)
		c.logger.Debugf(ctx, "%s\n", buf.String())
	} else {
		formatted, _ := json.MarshalIndent(m, "", "  ")
		c.logger.Debugf(ctx, "%s\n", formatted)
	}
	c.logger.Debugf(ctx, "<----- %s (response)\n", c.endpoint)

	return m
}
<commit_msg>client: discard resp.Body data before Close<commit_after>// Copyright 2017 The go-rollbar Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rollbar

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"time"

	"github.com/pkg/errors"
	api "github.com/zchee/go-rollbar/api/v1"
	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
)

type client struct {
	debugClient    *httpClient
	infoClient     *httpClient
	errorClient    *httpClient
	warnClient     *httpClient
	criticalClient *httpClient
}

type httpClient struct {
	token        string
	client       *http.Client
	endpoint     string
	debug        bool
	logger       Logger
	environment  string
	platform     string
	codeVersion  string
	serverHost   string
	serverRoot   string
	serverBranch string
}

var defaultHTTPClient = httpClient{
	client:      http.DefaultClient,
	endpoint:    api.DefaultEndpoint,
	logger:      nilLogger{},
	environment: "development",
	platform:    runtime.GOOS,
}

// New creates a new REST rollbar API client.
//
// The `token` is required, other optional parameters can be passed using the
// various `With...` functions.
func New(token string, options ...Option) Client {
	cl := defaultHTTPClient
	cl.token = token
	if debug, err := strconv.ParseBool(os.Getenv("ROLLBAR_DEBUG")); err == nil && debug {
		cl.debug = debug
	}

	for _, o := range options {
		o(&cl)
	}
	if _, ok := cl.logger.(nilLogger); cl.debug && ok {
		cl.logger = traceLogger{os.Stderr}
	}
	if cl.serverHost == "" {
		cl.serverHost, _ = os.Hostname()
	}

	return &client{
		debugClient:    &cl,
		infoClient:     &cl,
		errorClient:    &cl,
		warnClient:     &cl,
		criticalClient: &cl,
	}
}

// payload creates the rollbar payload data.
func (c *httpClient) payload(level Level, err error) *api.Payload {
	title := "<nil>"
	if err != nil {
		title = err.Error()
	}
	stack := CreateStack(3)

	data := &api.Data{
		Environment: c.environment,
		Body:        errorBody(err, stack),
		Level:       string(level),
		Timestamp:   time.Now().Unix(),
		Platform:    c.platform,
		Language:    language,
		Server: &api.Server{
			Host:   c.serverHost,
			Root:   c.serverRoot,
			Branch: c.serverBranch,
		},
		Fingerprint: stack.Fingerprint(),
		Title:       title,
		Notifier: &api.Notifier{
			Name:    Name,
			Version: Version,
		},
	}

	return &api.Payload{
		AccessToken: c.token,
		Data:        data,
	}
}

// post posts payload to rollbar.
func (c *httpClient) post(pctx context.Context, payload *api.Payload) (*api.Response, error) {
	if c.token == "" {
		return nil, errors.New("empty token")
	}

	data, err := json.Marshal(payload)
	if err != nil {
		return nil, errors.Wrap(err, "failed to encode payload")
	}
	if c.debug {
		out, _ := json.MarshalIndent(payload, "", "  ")
		c.logger.Debugf(pctx, string(out)+"\n")
	}

	req, err := http.NewRequest(http.MethodPost, c.endpoint, bytes.NewReader(data))
	if err != nil {
		return nil, errors.Wrap(err, "failed to create new POST request")
	}

	ctx, cancel := context.WithCancel(pctx)
	defer cancel()

	req.Header.Set("Content-Type", "application/json")
	resp, err := ctxhttp.Do(ctx, c.client, req)
	if err != nil {
		return nil, errors.Wrap(err, "failed to POST to rollbar")
	}
	defer func() {
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}()

	if resp.StatusCode != http.StatusOK {
		return nil, errors.Errorf("received response: %s", resp.Status)
	}

	return c.parseResponse(ctx, resp.Body), nil
}

// parseResponse parses the rollbar API response.
func (c *httpClient) parseResponse(ctx context.Context, rdr io.Reader) *api.Response {
	buf := new(bytes.Buffer)
	io.Copy(buf, rdr)

	c.logger.Debugf(ctx, "-----> %s (response)\n", c.endpoint)
	m := new(api.Response)
	if err := json.Unmarshal(buf.Bytes(), m); err != nil {
		c.logger.Debugf(ctx, "failed to unmarshal payload: %s\n", err)
		c.logger.Debugf(ctx, "%s\n", buf.String())
	} else {
		formatted, _ := json.MarshalIndent(m, "", "  ")
		c.logger.Debugf(ctx, "%s\n", formatted)
	}
	c.logger.Debugf(ctx, "<----- %s (response)\n", c.endpoint)

	return m
}
<|endoftext|>


<commit_before>// Package musicmanager implements a client for managing Google Play
// Music libraries.
package musicmanager

import (
	"fmt"
	"net/http"

	convert "github.com/lxr/go.google.musicmanager/internal/convert"
	mmdspb "github.com/lxr/go.google.musicmanager/internal/download_proto/service"
	mmldpb "github.com/lxr/go.google.musicmanager/internal/locker_proto/data"
	mmssjs "github.com/lxr/go.google.musicmanager/internal/session_json"
	mmudpb "github.com/lxr/go.google.musicmanager/internal/upload_proto/data"
	mmuspb "github.com/lxr/go.google.musicmanager/internal/upload_proto/service"
)

// Scope is the OAuth scope used by this API.
const Scope = "https://www.googleapis.com/auth/musicmanager"

// Client is a Music Manager client.
type Client struct {
	client *http.Client
	id     string
}

// NewClient creates a new Music Manager client with the given device ID
// and underlying HTTP client.  The supplied client must be capable of
// authenticating requests with the OAuth scope Scope.
func NewClient(client *http.Client, deviceID string) (*Client, error) {
	if client == nil {
		return nil, fmt.Errorf("client is nil")
	}
	if deviceID == "" {
		return nil, fmt.Errorf("device ID is empty")
	}
	return &Client{client, deviceID}, nil
}

// Register registers the client as a device with the given name under
// the user's Play Music library.  A client must be registered before
// the other method calls can succeed.  Re-registering a client can be
// used to change its name.  Note that there are limits to how many
// devices one account can have registered at a time, with how many
// accounts one device can be registered, and how many devices a user
// can deregister in a year, so be judicious in using this call.
func (c *Client) Register(name string) error {
	res, err := c.upAuth(&mmuspb.UpAuthRequest{
		UploaderId:   c.id,
		FriendlyName: name,
	})
	if err != nil {
		return err
	}
	if res != mmuspb.UploadResponse_OK {
		return RegisterError(res)
	}
	return nil
}

// ExportTrack returns a short-lived download URL for the given track,
// identified by its server ID.  Downloading the track from this URL
// requires no authentication.  The Content-Disposition and
// Content-Length headers of the response contain the name and size of
// the track respectively.
func (c *Client) ExportTrack(id string) (string, error) {
	res, err := c.getDownloadSession(&mmssjs.GetDownloadSessionRequest{
		XDeviceID: c.id,
		SongID:    id,
	})
	if err != nil {
		return "", err
	}
	return res.URL, nil
}

// ListTracks lists the user's tracks from least to most recently
// accessed.  It can optionally filter for purchased or promotional
// tracks only, or for tracks last modified after a given Unix timestamp
// (microsecond precision).  Long responses may be returned in chunks,
// in which case the PageToken field of the TrackList object should
// be given to a new ListTracks call.
func (c *Client) ListTracks(purchasedOnly bool, updatedMin int64, pageToken string) (*TrackList, error) {
	var exportType mmdspb.GetTracksToExportRequest_TracksToExportType
	switch purchasedOnly {
	case true:
		exportType = mmdspb.GetTracksToExportRequest_PURCHASED_AND_PROMOTIONAL
	case false:
		exportType = mmdspb.GetTracksToExportRequest_ALL
	}
	res, err := c.getTracksToExport(&mmdspb.GetTracksToExportRequest{
		ClientId:          c.id,
		ExportType:        exportType,
		UpdatedMin:        updatedMin,
		ContinuationToken: pageToken,
	})
	switch err := err.(type) {
	case nil:
		if res.Status != mmdspb.GetTracksToExportResponse_OK {
			return nil, ListError(res.Status)
		}
	case *RequestError:
		// The Google Play servers respond with 304 Not Modified
		// if no tracks have been modified after the updatedMin
		// timestamp.  This is not exactly an error condition,
		// so we break out of the switch rather than return the
		// error; the remaining function body then arranges to
		// return an empty TrackList.
		if err.Code == http.StatusNotModified {
			break
		}
		return nil, err
	default:
		return nil, err
	}
	trackList := new(TrackList)
	convert.Convert(trackList, res)
	trackList.PurchasedOnly = purchasedOnly
	return trackList, nil
}

// ImportTracks returns short-lived upload URLs for the given tracks.
// MP3 audio data can be PUT or POSTed to these URLs without
// authentication.  Individual tracks can fail, in which case errs[i]
// contains the reason why importing tracks[i] failed.
//
// The Title field of a track to be imported cannot be empty.  The
// server also uses the ClientId field to identify which tracks have
// already been uploaded, so trying to import a track with a client ID
// that has already been uploaded will fail.  Leaving the ClientId field
// empty appears to bypass this server-side check; however, the
// implementation of ImportTracks cannot handle two tracks in a batch
// with the same client ID, even if they are both empty.
//
// A response to a request sent to a URL returned from this function can
// report failure through means other than the status code.  Use
// CheckImportResponse to verify that the request succeeded.
func (c *Client) ImportTracks(tracks []*Track) (urls []string, errs []error) {
	// Construct and the client-ID-to-track-index mapping and the
	// initial metadata upload.
	cidm := make(map[string]int)
	trks := make([]*mmldpb.Track, len(tracks))
	errs = make([]error, len(tracks))
	for i, track := range tracks {
		if _, ok := cidm[track.ClientId]; ok {
			errs[i] = fmt.Errorf("trying to import two tracks with the same client-side ID")
			continue
		}
		trks[i] = new(mmldpb.Track)
		convert.Convert(trks[i], track)
		cidm[track.ClientId] = i
	}
	// Upload track metadata.
	res, err := c.uploadMetadata(&mmuspb.UploadMetadataRequest{
		UploaderId: c.id,
		Track:      trks,
	})
	if err != nil {
		for i := range errs {
			if errs[i] == nil {
				errs[i] = err
			}
		}
		return nil, errs
	}
	// Satisfy any requests for track samples and append the
	// responses to tres.
	tres := res.TrackSampleResponse
	if n := len(res.SignedChallengeInfo); n > 0 {
		spls := make([]*mmudpb.TrackSample, n)
		for i, sci := range res.SignedChallengeInfo {
			ci := sci.ChallengeInfo
			j := cidm[ci.ClientTrackId]
			var sample []byte
			if sf := tracks[j].SampleFunc; sf != nil {
				sample = sf(
					int(ci.StartMillis),
					int(ci.DurationMillis),
				)
			}
			if sample == nil {
				// A nil sample is different from an
				// empty one: the former results in an
				// invalid sample message.  So, if
				// Sampler leaves us with a nil sample,
				// replace it with an empty one.
				sample = make([]byte, 0)
			}
			spls[i] = &mmudpb.TrackSample{
				Track:               trks[j],
				Sample:              sample,
				SignedChallengeInfo: sci,
				SampleFormat:        mmldpb.Track_MP3,
			}
		}
		sres, err := c.uploadSample(&mmuspb.UploadSampleRequest{
			UploaderId:  c.id,
			TrackSample: spls,
		})
		if err != nil {
			for _, spl := range spls {
				errs[cidm[spl.Track.ClientId]] = err
			}
		} else {
			tres = append(tres, sres.TrackSampleResponse...)
		}
	}
	// Parse responses to track metadata and samples.  The result
	// is a map from track indices to their server IDs.
	sidm := make(map[int]string)
	for _, res := range tres {
		i := cidm[res.ClientTrackId]
		if res.ResponseCode != mmuspb.TrackSampleResponse_UPLOAD_REQUESTED {
			errs[i] = ImportError(res.ResponseCode)
			continue
		}
		sidm[i] = res.ServerTrackId
	}
	// Acquire upload sessions.
	urls = make([]string, len(tracks))
	for i, id := range sidm {
		trk := tracks[i]
		res, err := c.getUploadSession(&mmssjs.GetUploadSessionRequest{
			Name:         id,
			UploaderId:   c.id,
			ClientId:     trk.ClientId,
			ServerId:     id,
			TrackBitRate: trk.BitRate,
			SyncNow:      true,
			// BUG(lor): Client.ImportTracks does not
			// activate the upload progress tracker on
			// https://play.google.com/music.
		})
		if err != nil {
			errs[i] = err
			continue
		}
		if res.Error != nil {
			errs[i] = res.Error
			continue
		}
		urls[i] = res.Transfers[0].PutUrl
	}
	return urls, errs
}

// CheckImportResponse checks the response to an HTTP request sent to a
// URL returned by Client.ImportTracks.  It returns the server ID of the
// track to which resp is a response, or an error if the response is
// erroneous in some way.  The error will be of type *RequestError if
// the response has a non-2xx status code.  CheckImportResponse closes
// resp.Body.
func CheckImportResponse(resp *http.Response) (string, error) {
	res := new(mmssjs.GetUploadSessionResponse)
	err := parseResponse(resp, res)
	switch {
	case err != nil:
		return "", err
	case res.Error != nil:
		return "", res.Error
	default:
		return res.Transfers[0].Name, nil
	}
}
<commit_msg>Document the "non-MAC-address-like" device ID download bug<commit_after>// Package musicmanager implements a client for managing Google Play
// Music libraries.
package musicmanager

import (
	"fmt"
	"net/http"

	convert "github.com/lxr/go.google.musicmanager/internal/convert"
	mmdspb "github.com/lxr/go.google.musicmanager/internal/download_proto/service"
	mmldpb "github.com/lxr/go.google.musicmanager/internal/locker_proto/data"
	mmssjs "github.com/lxr/go.google.musicmanager/internal/session_json"
	mmudpb "github.com/lxr/go.google.musicmanager/internal/upload_proto/data"
	mmuspb "github.com/lxr/go.google.musicmanager/internal/upload_proto/service"
)

// Scope is the OAuth scope used by this API.
const Scope = "https://www.googleapis.com/auth/musicmanager"

// Client is a Music Manager client.
type Client struct {
	client *http.Client
	id     string
}

// NewClient creates a new Music Manager client with the given device ID
// and underlying HTTP client.  The supplied client must be capable of
// authenticating requests with the OAuth scope Scope.
func NewClient(client *http.Client, deviceID string) (*Client, error) {
	if client == nil {
		return nil, fmt.Errorf("client is nil")
	}
	if deviceID == "" {
		return nil, fmt.Errorf("device ID is empty")
	}
	return &Client{client, deviceID}, nil
}

// BUG(lor): Exporting tracks has been known to fail if the device ID of
// the client is not sufficiently "MAC address-like".  The exact
// threshold is unknown; perhaps the server only looks for a colon in
// the ID.

// Register registers the client as a device with the given name under
// the user's Play Music library.  A client must be registered before
// the other method calls can succeed.  Re-registering a client can be
// used to change its name.  Note that there are limits to how many
// devices one account can have registered at a time, with how many
// accounts one device can be registered, and how many devices a user
// can deregister in a year, so be judicious in using this call.
func (c *Client) Register(name string) error {
	res, err := c.upAuth(&mmuspb.UpAuthRequest{
		UploaderId:   c.id,
		FriendlyName: name,
	})
	if err != nil {
		return err
	}
	if res != mmuspb.UploadResponse_OK {
		return RegisterError(res)
	}
	return nil
}

// ExportTrack returns a short-lived download URL for the given track,
// identified by its server ID.  Downloading the track from this URL
// requires no authentication.  The Content-Disposition and
// Content-Length headers of the response contain the name and size of
// the track respectively.
func (c *Client) ExportTrack(id string) (string, error) {
	res, err := c.getDownloadSession(&mmssjs.GetDownloadSessionRequest{
		XDeviceID: c.id,
		SongID:    id,
	})
	if err != nil {
		return "", err
	}
	return res.URL, nil
}

// ListTracks lists the user's tracks from least to most recently
// accessed.  It can optionally filter for purchased or promotional
// tracks only, or for tracks last modified after a given Unix timestamp
// (microsecond precision).  Long responses may be returned in chunks,
// in which case the PageToken field of the TrackList object should
// be given to a new ListTracks call.
func (c *Client) ListTracks(purchasedOnly bool, updatedMin int64, pageToken string) (*TrackList, error) {
	var exportType mmdspb.GetTracksToExportRequest_TracksToExportType
	switch purchasedOnly {
	case true:
		exportType = mmdspb.GetTracksToExportRequest_PURCHASED_AND_PROMOTIONAL
	case false:
		exportType = mmdspb.GetTracksToExportRequest_ALL
	}
	res, err := c.getTracksToExport(&mmdspb.GetTracksToExportRequest{
		ClientId:          c.id,
		ExportType:        exportType,
		UpdatedMin:        updatedMin,
		ContinuationToken: pageToken,
	})
	switch err := err.(type) {
	case nil:
		if res.Status != mmdspb.GetTracksToExportResponse_OK {
			return nil, ListError(res.Status)
		}
	case *RequestError:
		// The Google Play servers respond with 304 Not Modified
		// if no tracks have been modified after the updatedMin
		// timestamp.  This is not exactly an error condition,
		// so we break out of the switch rather than return the
		// error; the remaining function body then arranges to
		// return an empty TrackList.
		if err.Code == http.StatusNotModified {
			break
		}
		return nil, err
	default:
		return nil, err
	}
	trackList := new(TrackList)
	convert.Convert(trackList, res)
	trackList.PurchasedOnly = purchasedOnly
	return trackList, nil
}

// ImportTracks returns short-lived upload URLs for the given tracks.
// MP3 audio data can be PUT or POSTed to these URLs without
// authentication.  Individual tracks can fail, in which case errs[i]
// contains the reason why importing tracks[i] failed.
//
// The Title field of a track to be imported cannot be empty.  The
// server also uses the ClientId field to identify which tracks have
// already been uploaded, so trying to import a track with a client ID
// that has already been uploaded will fail.  Leaving the ClientId field
// empty appears to bypass this server-side check; however, the
// implementation of ImportTracks cannot handle two tracks in a batch
// with the same client ID, even if they are both empty.
//
// A response to a request sent to a URL returned from this function can
// report failure through means other than the status code.  Use
// CheckImportResponse to verify that the request succeeded.
func (c *Client) ImportTracks(tracks []*Track) (urls []string, errs []error) {
	// Construct and the client-ID-to-track-index mapping and the
	// initial metadata upload.
	cidm := make(map[string]int)
	trks := make([]*mmldpb.Track, len(tracks))
	errs = make([]error, len(tracks))
	for i, track := range tracks {
		if _, ok := cidm[track.ClientId]; ok {
			errs[i] = fmt.Errorf("trying to import two tracks with the same client-side ID")
			continue
		}
		trks[i] = new(mmldpb.Track)
		convert.Convert(trks[i], track)
		cidm[track.ClientId] = i
	}
	// Upload track metadata.
	res, err := c.uploadMetadata(&mmuspb.UploadMetadataRequest{
		UploaderId: c.id,
		Track:      trks,
	})
	if err != nil {
		for i := range errs {
			if errs[i] == nil {
				errs[i] = err
			}
		}
		return nil, errs
	}
	// Satisfy any requests for track samples and append the
	// responses to tres.
	tres := res.TrackSampleResponse
	if n := len(res.SignedChallengeInfo); n > 0 {
		spls := make([]*mmudpb.TrackSample, n)
		for i, sci := range res.SignedChallengeInfo {
			ci := sci.ChallengeInfo
			j := cidm[ci.ClientTrackId]
			var sample []byte
			if sf := tracks[j].SampleFunc; sf != nil {
				sample = sf(
					int(ci.StartMillis),
					int(ci.DurationMillis),
				)
			}
			if sample == nil {
				// A nil sample is different from an
				// empty one: the former results in an
				// invalid sample message.  So, if
				// Sampler leaves us with a nil sample,
				// replace it with an empty one.
				sample = make([]byte, 0)
			}
			spls[i] = &mmudpb.TrackSample{
				Track:               trks[j],
				Sample:              sample,
				SignedChallengeInfo: sci,
				SampleFormat:        mmldpb.Track_MP3,
			}
		}
		sres, err := c.uploadSample(&mmuspb.UploadSampleRequest{
			UploaderId:  c.id,
			TrackSample: spls,
		})
		if err != nil {
			for _, spl := range spls {
				errs[cidm[spl.Track.ClientId]] = err
			}
		} else {
			tres = append(tres, sres.TrackSampleResponse...)
		}
	}
	// Parse responses to track metadata and samples.  The result
	// is a map from track indices to their server IDs.
	sidm := make(map[int]string)
	for _, res := range tres {
		i := cidm[res.ClientTrackId]
		if res.ResponseCode != mmuspb.TrackSampleResponse_UPLOAD_REQUESTED {
			errs[i] = ImportError(res.ResponseCode)
			continue
		}
		sidm[i] = res.ServerTrackId
	}
	// Acquire upload sessions.
	urls = make([]string, len(tracks))
	for i, id := range sidm {
		trk := tracks[i]
		res, err := c.getUploadSession(&mmssjs.GetUploadSessionRequest{
			Name:         id,
			UploaderId:   c.id,
			ClientId:     trk.ClientId,
			ServerId:     id,
			TrackBitRate: trk.BitRate,
			SyncNow:      true,
			// BUG(lor): Client.ImportTracks does not
			// activate the upload progress tracker on
			// https://play.google.com/music.
		})
		if err != nil {
			errs[i] = err
			continue
		}
		if res.Error != nil {
			errs[i] = res.Error
			continue
		}
		urls[i] = res.Transfers[0].PutUrl
	}
	return urls, errs
}

// CheckImportResponse checks the response to an HTTP request sent to a
// URL returned by Client.ImportTracks.  It returns the server ID of the
// track to which resp is a response, or an error if the response is
// erroneous in some way.  The error will be of type *RequestError if
// the response has a non-2xx status code.  CheckImportResponse closes
// resp.Body.
func CheckImportResponse(resp *http.Response) (string, error) {
	res := new(mmssjs.GetUploadSessionResponse)
	err := parseResponse(resp, res)
	switch {
	case err != nil:
		return "", err
	case res.Error != nil:
		return "", res.Error
	default:
		return res.Transfers[0].Name, nil
	}
}
<|endoftext|>


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

package rpcmq

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/streadway/amqp"
)

// A Client is an RPC client, which is used to invoke remote procedures.
type Client struct {
	msgsName     string
	repliesName  string
	exchangeName string
	exchangeKind string
	ac           *amqpClient
	msgsQueue    amqp.Queue
	repliesQueue amqp.Queue
	mandatory    bool
	deliveries   <-chan amqp.Delivery
	results      chan Result

	// DeliveryMode allows to configure the delivery mode followed by the
	// broker. The default mode is Persistent.
	DeliveryMode DeliveryMode

	// TLSConfig allows to configure the TLS parameters used to connect to
	// the broker via amqps.
	TLSConfig *tls.Config
}

// A Result contains the data returned by the invoked procedure or an error
// message, in case that it finished with error. The UUID allows to link the
// result with the procedure call.
type Result struct {
	UUID string
	Data []byte
	Err  string
}

type rpcMsg struct {
	Method string
	Data   []byte
}

// NewClient returns a reference to a Client object. The paremeter uri is the
// network address of the broker and msgsQueue/repliesQueue are the names of
// queues that will be created to exchange the messages between clients and
// servers. On the other hand, the parameters exchange and kind determine the
// type of exchange that will be created. In fanout mode the queue name is
// ignored, so each queue has its own unique id.
func NewClient(uri, msgsQueue, repliesQueue, exchange, kind string) *Client {
	if kind == "fanout" {
		msgsQueue = "" // in fanout mode queue names must be unique
	}
	c := &Client{
		msgsName:     msgsQueue,
		repliesName:  repliesQueue,
		exchangeName: exchange,
		exchangeKind: kind,
		ac:           newAmqpClient(uri),
		results:      make(chan Result),
		DeliveryMode: Persistent,
	}
	c.ac.setupFunc = c.setup
	return c
}

// Init initializes the Client object. It establishes the connection with the
// broker, creating a channel and the queues that will be used under the hood.
func (c *Client) Init() error {
	c.ac.tlsConfig = c.TLSConfig
	if err := c.ac.init(); err != nil {
		return err
	}
	return c.setup()
}

func (c *Client) setup() error {
	err := c.ac.channel.ExchangeDeclare(
		c.exchangeName, // name
		c.exchangeKind, // kind
		true,           // durable
		false,          // autoDelete
		false,          // internal
		false,          // noWait
		nil,            // args
	)
	if err != nil {
		return fmt.Errorf("ExchangeDeclare: %v", err)
	}

	if c.exchangeKind != "fanout" {
		// We should create the queue only in non-fanout mode
		c.msgsQueue, err = c.ac.channel.QueueDeclare(
			c.msgsName, // name
			true,       // durable
			false,      // autoDelete
			false,      // exclusive
			false,      // noWait
			nil,        // args
		)
		if err != nil {
			return fmt.Errorf("QueueDeclare: %v", err)
		}

		err = c.ac.channel.QueueBind(
			c.msgsQueue.Name, // name
			c.msgsQueue.Name, // key
			c.exchangeName,   // exchange
			false,            // noWait
			nil,              // args
		)
		if err != nil {
			return fmt.Errorf("QueueBind: %v", err)
		}
	}

	c.repliesQueue, err = c.ac.channel.QueueDeclare(
		c.repliesName, // name
		true,          // durable
		false,         // autoDelete
		false,         // exclusive
		false,         // noWait
		nil,           // args
	)
	if err != nil {
		return fmt.Errorf("QueueDeclare: %v", err)
	}

	c.ac.consumerTag, err = uuid()
	if err != nil {
		return fmt.Errorf("UUID: %v", err)
	}

	c.deliveries, err = c.ac.channel.Consume(
		c.repliesQueue.Name, // name
		c.ac.consumerTag,    // consumer
		false,               // autoAck
		false,               // exclusive
		false,               // noLocal
		false,               // noWait
		nil,                 // args
	)
	if err != nil {
		return fmt.Errorf("Queue Consume: %v", err)
	}

	go c.getDeliveries()

	return nil
}

func (c *Client) getDeliveries() {
	for d := range c.deliveries {
		if d.CorrelationId == "" {
			d.Nack(false, false) // drop message
			logf("dropped message: %+v", d)
			continue
		}
		var r Result
		if err := json.Unmarshal(d.Body, &r); err != nil {
			d.Nack(false, false) // drop message
			logf("dropped message: %+v", d)
			continue
		}
		c.results <- r
		d.Ack(false)
	}
	c.ac.done <- true
}

// Shutdown shuts down the client gracefully. Using this method will ensure
// that all replies sent by the RPC servers to the client will be received by
// the latter.
func (c *Client) Shutdown() {
	c.ac.shutdown()
}

// Call invokes the remote procedure specified by the parameter method, being
// the parameter data the input passed to it. On the other hand, ttl is the
// time that this task will remain in the queue before being considered dead.
// The returned id can be used to identify the result corresponding to each
// invokation. If ttl is 0, the message will not expire.
func (c *Client) Call(method string, data []byte, ttl time.Duration) (id string, err error) {
	id, err = uuid()
	if err != nil {
		return "", fmt.Errorf("UUID: %v", err)
	}

	msg := &rpcMsg{
		Method: method,
		Data:   data,
	}
	body, err := json.Marshal(msg)
	if err != nil {
		return "", fmt.Errorf("Marshal: %v", err)
	}

	expiration := ""
	if ttl > 0 {
		expiration = fmt.Sprintf("%d", int64(ttl.Seconds()*1000))
	}
	c.mandatory = true
	if c.exchangeKind == "fanout" {
		// In fanout mode the routing key is not really used, so using
		// the mandatory flag does not make sense
		c.mandatory = false
	}

	// guarantee that the received ack/nack corresponds with this publishing
	c.ac.mu.Lock()
	defer c.ac.mu.Unlock()

	err = c.ac.channel.Publish(
		c.exchangeName, // exchange
		c.msgsName,     // key
		c.mandatory,    // mandatory
		false,          // immediate
		amqp.Publishing{ // msg
			CorrelationId: id,
			ReplyTo:       c.repliesQueue.Name,
			ContentType:   "application/json",
			Body:          body,
			DeliveryMode:  uint8(c.DeliveryMode),
			Expiration:    expiration,
		},
	)
	if err != nil {
		return "", err
	}

	select {
	case _, ok := <-c.ac.acks:
		if ok {
			return id, nil
		}
	case tag, ok := <-c.ac.nacks:
		if ok {
			logf("nack recived (%v)", tag)
			return "", errors.New("nack received")
		}
	}

	logf("missing ack/nack")
	return "", errors.New("missing ack/nack")
}

// Results returns a channel used to receive the results returned by the
// invoked procedures.
func (c *Client) Results() <-chan Result {
	return (<-chan Result)(c.results)
}
<commit_msg>Set mandatory value in NewClient<commit_after>// Copyright 2015 The rpcmq Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpcmq

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/streadway/amqp"
)

// A Client is an RPC client, which is used to invoke remote procedures.
type Client struct {
	msgsName     string
	repliesName  string
	exchangeName string
	exchangeKind string
	ac           *amqpClient
	msgsQueue    amqp.Queue
	repliesQueue amqp.Queue
	mandatory    bool
	deliveries   <-chan amqp.Delivery
	results      chan Result

	// DeliveryMode allows to configure the delivery mode followed by the
	// broker. The default mode is Persistent.
	DeliveryMode DeliveryMode

	// TLSConfig allows to configure the TLS parameters used to connect to
	// the broker via amqps.
	TLSConfig *tls.Config
}

// A Result contains the data returned by the invoked procedure or an error
// message, in case that it finished with error. The UUID allows to link the
// result with the procedure call.
type Result struct {
	UUID string
	Data []byte
	Err  string
}

type rpcMsg struct {
	Method string
	Data   []byte
}

// NewClient returns a reference to a Client object. The paremeter uri is the
// network address of the broker and msgsQueue/repliesQueue are the names of
// queues that will be created to exchange the messages between clients and
// servers. On the other hand, the parameters exchange and kind determine the
// type of exchange that will be created. In fanout mode the queue name is
// ignored, so each queue has its own unique id.
func NewClient(uri, msgsQueue, repliesQueue, exchange, kind string) *Client {
	if kind == "fanout" {
		msgsQueue = "" // in fanout mode queue names must be unique
	}
	c := &Client{
		msgsName:     msgsQueue,
		repliesName:  repliesQueue,
		exchangeName: exchange,
		exchangeKind: kind,
		ac:           newAmqpClient(uri),
		results:      make(chan Result),
		DeliveryMode: Persistent,
	}

	// In fanout mode the routing key is not really used, so using
	// the mandatory flag does not make sense
	if kind == "fanout" {
		c.mandatory = false
	} else {
		c.mandatory = true
	}

	c.ac.setupFunc = c.setup
	return c
}

// Init initializes the Client object. It establishes the connection with the
// broker, creating a channel and the queues that will be used under the hood.
func (c *Client) Init() error {
	c.ac.tlsConfig = c.TLSConfig
	if err := c.ac.init(); err != nil {
		return err
	}
	return c.setup()
}

func (c *Client) setup() error {
	err := c.ac.channel.ExchangeDeclare(
		c.exchangeName, // name
		c.exchangeKind, // kind
		true,           // durable
		false,          // autoDelete
		false,          // internal
		false,          // noWait
		nil,            // args
	)
	if err != nil {
		return fmt.Errorf("ExchangeDeclare: %v", err)
	}

	if c.exchangeKind != "fanout" {
		// We should create the queue only in non-fanout mode
		c.msgsQueue, err = c.ac.channel.QueueDeclare(
			c.msgsName, // name
			true,       // durable
			false,      // autoDelete
			false,      // exclusive
			false,      // noWait
			nil,        // args
		)
		if err != nil {
			return fmt.Errorf("QueueDeclare: %v", err)
		}

		err = c.ac.channel.QueueBind(
			c.msgsQueue.Name, // name
			c.msgsQueue.Name, // key
			c.exchangeName,   // exchange
			false,            // noWait
			nil,              // args
		)
		if err != nil {
			return fmt.Errorf("QueueBind: %v", err)
		}
	}

	c.repliesQueue, err = c.ac.channel.QueueDeclare(
		c.repliesName, // name
		true,          // durable
		false,         // autoDelete
		false,         // exclusive
		false,         // noWait
		nil,           // args
	)
	if err != nil {
		return fmt.Errorf("QueueDeclare: %v", err)
	}

	c.ac.consumerTag, err = uuid()
	if err != nil {
		return fmt.Errorf("UUID: %v", err)
	}

	c.deliveries, err = c.ac.channel.Consume(
		c.repliesQueue.Name, // name
		c.ac.consumerTag,    // consumer
		false,               // autoAck
		false,               // exclusive
		false,               // noLocal
		false,               // noWait
		nil,                 // args
	)
	if err != nil {
		return fmt.Errorf("Queue Consume: %v", err)
	}

	go c.getDeliveries()

	return nil
}

func (c *Client) getDeliveries() {
	for d := range c.deliveries {
		if d.CorrelationId == "" {
			d.Nack(false, false) // drop message
			logf("dropped message: %+v", d)
			continue
		}
		var r Result
		if err := json.Unmarshal(d.Body, &r); err != nil {
			d.Nack(false, false) // drop message
			logf("dropped message: %+v", d)
			continue
		}
		c.results <- r
		d.Ack(false)
	}
	c.ac.done <- true
}

// Shutdown shuts down the client gracefully. Using this method will ensure
// that all replies sent by the RPC servers to the client will be received by
// the latter.
func (c *Client) Shutdown() {
	c.ac.shutdown()
}

// Call invokes the remote procedure specified by the parameter method, being
// the parameter data the input passed to it. On the other hand, ttl is the
// time that this task will remain in the queue before being considered dead.
// The returned id can be used to identify the result corresponding to each
// invokation. If ttl is 0, the message will not expire.
func (c *Client) Call(method string, data []byte, ttl time.Duration) (id string, err error) {
	id, err = uuid()
	if err != nil {
		return "", fmt.Errorf("UUID: %v", err)
	}

	msg := &rpcMsg{
		Method: method,
		Data:   data,
	}
	body, err := json.Marshal(msg)
	if err != nil {
		return "", fmt.Errorf("Marshal: %v", err)
	}

	expiration := ""
	if ttl > 0 {
		expiration = fmt.Sprintf("%d", int64(ttl.Seconds()*1000))
	}

	// guarantee that the received ack/nack corresponds with this publishing
	c.ac.mu.Lock()
	defer c.ac.mu.Unlock()

	err = c.ac.channel.Publish(
		c.exchangeName, // exchange
		c.msgsName,     // key
		c.mandatory,    // mandatory
		false,          // immediate
		amqp.Publishing{ // msg
			CorrelationId: id,
			ReplyTo:       c.repliesQueue.Name,
			ContentType:   "application/json",
			Body:          body,
			DeliveryMode:  uint8(c.DeliveryMode),
			Expiration:    expiration,
		},
	)
	if err != nil {
		return "", err
	}

	select {
	case _, ok := <-c.ac.acks:
		if ok {
			return id, nil
		}
	case tag, ok := <-c.ac.nacks:
		if ok {
			logf("nack recived (%v)", tag)
			return "", errors.New("nack received")
		}
	}

	logf("missing ack/nack")
	return "", errors.New("missing ack/nack")
}

// Results returns a channel used to receive the results returned by the
// invoked procedures.
func (c *Client) Results() <-chan Result {
	return (<-chan Result)(c.results)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

const MSG_BUFFER int = 10

const HELP_TEXT string = `-> Available commands:
   /about
   /exit
   /help
   /list
   /nick $NAME
   /whois $NAME
`

const ABOUT_TEXT string = `-> ssh-chat is made by @shazow.

   It is a custom ssh server built in Go to serve a chat experience
   instead of a shell.

   Source: https://github.com/shazow/ssh-chat

   For more, visit shazow.net or follow at twitter.com/shazow
`

type Client struct {
	Server        *Server
	Conn          *ssh.ServerConn
	Msg           chan string
	Name          string
	Op            bool
	ready         chan struct{}
	term          *terminal.Terminal
	termWidth     int
	termHeight    int
	silencedUntil time.Time
}

func NewClient(server *Server, conn *ssh.ServerConn) *Client {
	return &Client{
		Server: server,
		Conn:   conn,
		Name:   conn.User(),
		Msg:    make(chan string, MSG_BUFFER),
		ready:  make(chan struct{}, 1),
	}
}

func (c *Client) Write(msg string) {
	c.term.Write([]byte(msg + "\r\n"))
}

func (c *Client) WriteLines(msg []string) {
	for _, line := range msg {
		c.Write(line)
	}
}

func (c *Client) IsSilenced() bool {
	return c.silencedUntil.After(time.Now())
}

func (c *Client) Silence(d time.Duration) {
	c.silencedUntil = time.Now().Add(d)
}

func (c *Client) Resize(width int, height int) error {
	err := c.term.SetSize(width, height)
	if err != nil {
		logger.Errorf("Resize failed: %dx%d", width, height)
		return err
	}
	c.termWidth, c.termHeight = width, height
	return nil
}

func (c *Client) Rename(name string) {
	c.Name = name
	c.term.SetPrompt(fmt.Sprintf("[%s] ", name))
}

func (c *Client) Fingerprint() string {
	return c.Conn.Permissions.Extensions["fingerprint"]
}

func (c *Client) handleShell(channel ssh.Channel) {
	defer channel.Close()

	// FIXME: This shouldn't live here, need to restructure the call chaining.
	c.Server.Add(c)
	go func() {
		// Block until done, then remove.
		c.Conn.Wait()
		c.Server.Remove(c)
	}()

	go func() {
		for msg := range c.Msg {
			c.Write(msg)
		}
	}()

	for {
		line, err := c.term.ReadLine()
		if err != nil {
			break
		}

		parts := strings.SplitN(line, " ", 3)
		isCmd := strings.HasPrefix(parts[0], "/")

		if isCmd {
			// TODO: Factor this out.
			switch parts[0] {
			case "/exit":
				channel.Close()
			case "/help":
				c.WriteLines(strings.Split(HELP_TEXT, "\n"))
			case "/about":
				c.WriteLines(strings.Split(ABOUT_TEXT, "\n"))
			case "/me":
				me := strings.TrimLeft(line, "/me")
				if me == "" {
					me = " is at a loss for words."
				}
				msg := fmt.Sprintf("** %s%s", c.Name, me)
				if c.IsSilenced() {
					c.Msg <- fmt.Sprintf("-> Message rejected, silenced.")
				} else {
					c.Server.Broadcast(msg[:200], nil)
				}
			case "/nick":
				if len(parts) == 2 {
					c.Server.Rename(c, parts[1])
				} else {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /nick $NAME")
				}
			case "/whois":
				if len(parts) == 2 {
					client := c.Server.Who(parts[1])
					if client != nil {
						version := client.Conn.ClientVersion()
						if len(version) > 100 {
							version = []byte("Evil Jerk with a superlong string")
						}
						c.Msg <- fmt.Sprintf("-> %s is %s via %s", client.Name, client.Fingerprint(), version)
					} else {
						c.Msg <- fmt.Sprintf("-> No such name: %s", parts[1])
					}
				} else {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /whois $NAME")
				}
			case "/list":
				names := c.Server.List(nil)
				c.Msg <- fmt.Sprintf("-> %d connected: %s", len(names), strings.Join(names, ", "))
			case "/ban":
				if !c.Server.IsOp(c) {
					c.Msg <- fmt.Sprintf("-> You're not an admin.")
				} else if len(parts) != 2 {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /ban $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.Msg <- fmt.Sprintf("-> No such name: %s", parts[1])
					} else {
						fingerprint := client.Fingerprint()
						client.Write(fmt.Sprintf("-> Banned by %s.", c.Name))
						c.Server.Ban(fingerprint, nil)
						client.Conn.Close()
						c.Server.Broadcast(fmt.Sprintf("* %s was banned by %s", parts[1], c.Name), nil)
					}
				}
			case "/op":
				if !c.Server.IsOp(c) {
					c.Msg <- fmt.Sprintf("-> You're not an admin.")
				} else if len(parts) != 2 {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /op $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.Msg <- fmt.Sprintf("-> No such name: %s", parts[1])
					} else {
						fingerprint := client.Fingerprint()
						client.Write(fmt.Sprintf("-> Made op by %s.", c.Name))
						c.Server.Op(fingerprint)
					}
				}
			case "/silence":
				if !c.Server.IsOp(c) {
					c.Msg <- fmt.Sprintf("-> You're not an admin.")
				} else if len(parts) < 2 {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /silence $NAME")
				} else {
					duration := time.Duration(5) * time.Minute
					if len(parts) >= 3 {
						parsedDuration, err := time.ParseDuration(parts[2])
						if err == nil {
							duration = parsedDuration
						}
					}
					client := c.Server.Who(parts[1])
					if client == nil {
						c.Msg <- fmt.Sprintf("-> No such name: %s", parts[1])
					} else {
						client.Silence(duration)
						client.Write(fmt.Sprintf("-> Silenced for %s by %s.", duration, c.Name))
					}
				}
			default:
				c.Msg <- fmt.Sprintf("-> Invalid command: %s", line)
			}
			continue
		}

		msg := fmt.Sprintf("%s: %s", c.Name, line)
		if c.IsSilenced() || len(msg) > 1000 {
			c.Msg <- fmt.Sprintf("-> Message rejected.")
			continue
		}
		c.Server.Broadcast(msg, c)
	}

}

func (c *Client) handleChannels(channels <-chan ssh.NewChannel) {
	prompt := fmt.Sprintf("[%s] ", c.Name)

	hasShell := false

	for ch := range channels {
		if t := ch.ChannelType(); t != "session" {
			ch.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
			continue
		}

		channel, requests, err := ch.Accept()
		if err != nil {
			logger.Errorf("Could not accept channel: %v", err)
			continue
		}
		defer channel.Close()

		c.term = terminal.NewTerminal(channel, prompt)
		for req := range requests {
			var width, height int
			var ok bool

			switch req.Type {
			case "shell":
				if c.term != nil && !hasShell {
					go c.handleShell(channel)
					ok = true
					hasShell = true
				}
			case "pty-req":
				width, height, ok = parsePtyRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			case "window-change":
				width, height, ok = parseWinchRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			}

			if req.WantReply {
				req.Reply(ok, nil)
			}
		}
	}
}
<commit_msg>Fixing len check.<commit_after>package main

import (
	"fmt"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

const MSG_BUFFER int = 10

const HELP_TEXT string = `-> Available commands:
   /about
   /exit
   /help
   /list
   /nick $NAME
   /whois $NAME
`

const ABOUT_TEXT string = `-> ssh-chat is made by @shazow.

   It is a custom ssh server built in Go to serve a chat experience
   instead of a shell.

   Source: https://github.com/shazow/ssh-chat

   For more, visit shazow.net or follow at twitter.com/shazow
`

type Client struct {
	Server        *Server
	Conn          *ssh.ServerConn
	Msg           chan string
	Name          string
	Op            bool
	ready         chan struct{}
	term          *terminal.Terminal
	termWidth     int
	termHeight    int
	silencedUntil time.Time
}

func NewClient(server *Server, conn *ssh.ServerConn) *Client {
	return &Client{
		Server: server,
		Conn:   conn,
		Name:   conn.User(),
		Msg:    make(chan string, MSG_BUFFER),
		ready:  make(chan struct{}, 1),
	}
}

func (c *Client) Write(msg string) {
	c.term.Write([]byte(msg + "\r\n"))
}

func (c *Client) WriteLines(msg []string) {
	for _, line := range msg {
		c.Write(line)
	}
}

func (c *Client) IsSilenced() bool {
	return c.silencedUntil.After(time.Now())
}

func (c *Client) Silence(d time.Duration) {
	c.silencedUntil = time.Now().Add(d)
}

func (c *Client) Resize(width int, height int) error {
	err := c.term.SetSize(width, height)
	if err != nil {
		logger.Errorf("Resize failed: %dx%d", width, height)
		return err
	}
	c.termWidth, c.termHeight = width, height
	return nil
}

func (c *Client) Rename(name string) {
	c.Name = name
	c.term.SetPrompt(fmt.Sprintf("[%s] ", name))
}

func (c *Client) Fingerprint() string {
	return c.Conn.Permissions.Extensions["fingerprint"]
}

func (c *Client) handleShell(channel ssh.Channel) {
	defer channel.Close()

	// FIXME: This shouldn't live here, need to restructure the call chaining.
	c.Server.Add(c)
	go func() {
		// Block until done, then remove.
		c.Conn.Wait()
		c.Server.Remove(c)
	}()

	go func() {
		for msg := range c.Msg {
			c.Write(msg)
		}
	}()

	for {
		line, err := c.term.ReadLine()
		if err != nil {
			break
		}

		parts := strings.SplitN(line, " ", 3)
		isCmd := strings.HasPrefix(parts[0], "/")

		if isCmd {
			// TODO: Factor this out.
			switch parts[0] {
			case "/exit":
				channel.Close()
			case "/help":
				c.WriteLines(strings.Split(HELP_TEXT, "\n"))
			case "/about":
				c.WriteLines(strings.Split(ABOUT_TEXT, "\n"))
			case "/me":
				me := strings.TrimLeft(line, "/me")
				if me == "" {
					me = " is at a loss for words."
				}
				msg := fmt.Sprintf("** %s%s", c.Name, me)
				if c.IsSilenced() || len(msg) > 1000 {
					c.Msg <- fmt.Sprintf("-> Message rejected.")
				} else {
					c.Server.Broadcast(msg, nil)
				}
			case "/nick":
				if len(parts) == 2 {
					c.Server.Rename(c, parts[1])
				} else {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /nick $NAME")
				}
			case "/whois":
				if len(parts) == 2 {
					client := c.Server.Who(parts[1])
					if client != nil {
						version := client.Conn.ClientVersion()
						if len(version) > 100 {
							version = []byte("Evil Jerk with a superlong string")
						}
						c.Msg <- fmt.Sprintf("-> %s is %s via %s", client.Name, client.Fingerprint(), version)
					} else {
						c.Msg <- fmt.Sprintf("-> No such name: %s", parts[1])
					}
				} else {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /whois $NAME")
				}
			case "/list":
				names := c.Server.List(nil)
				c.Msg <- fmt.Sprintf("-> %d connected: %s", len(names), strings.Join(names, ", "))
			case "/ban":
				if !c.Server.IsOp(c) {
					c.Msg <- fmt.Sprintf("-> You're not an admin.")
				} else if len(parts) != 2 {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /ban $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.Msg <- fmt.Sprintf("-> No such name: %s", parts[1])
					} else {
						fingerprint := client.Fingerprint()
						client.Write(fmt.Sprintf("-> Banned by %s.", c.Name))
						c.Server.Ban(fingerprint, nil)
						client.Conn.Close()
						c.Server.Broadcast(fmt.Sprintf("* %s was banned by %s", parts[1], c.Name), nil)
					}
				}
			case "/op":
				if !c.Server.IsOp(c) {
					c.Msg <- fmt.Sprintf("-> You're not an admin.")
				} else if len(parts) != 2 {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /op $NAME")
				} else {
					client := c.Server.Who(parts[1])
					if client == nil {
						c.Msg <- fmt.Sprintf("-> No such name: %s", parts[1])
					} else {
						fingerprint := client.Fingerprint()
						client.Write(fmt.Sprintf("-> Made op by %s.", c.Name))
						c.Server.Op(fingerprint)
					}
				}
			case "/silence":
				if !c.Server.IsOp(c) {
					c.Msg <- fmt.Sprintf("-> You're not an admin.")
				} else if len(parts) < 2 {
					c.Msg <- fmt.Sprintf("-> Missing $NAME from: /silence $NAME")
				} else {
					duration := time.Duration(5) * time.Minute
					if len(parts) >= 3 {
						parsedDuration, err := time.ParseDuration(parts[2])
						if err == nil {
							duration = parsedDuration
						}
					}
					client := c.Server.Who(parts[1])
					if client == nil {
						c.Msg <- fmt.Sprintf("-> No such name: %s", parts[1])
					} else {
						client.Silence(duration)
						client.Write(fmt.Sprintf("-> Silenced for %s by %s.", duration, c.Name))
					}
				}
			default:
				c.Msg <- fmt.Sprintf("-> Invalid command: %s", line)
			}
			continue
		}

		msg := fmt.Sprintf("%s: %s", c.Name, line)
		if c.IsSilenced() || len(msg) > 1000 {
			c.Msg <- fmt.Sprintf("-> Message rejected.")
			continue
		}
		c.Server.Broadcast(msg, c)
	}

}

func (c *Client) handleChannels(channels <-chan ssh.NewChannel) {
	prompt := fmt.Sprintf("[%s] ", c.Name)

	hasShell := false

	for ch := range channels {
		if t := ch.ChannelType(); t != "session" {
			ch.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
			continue
		}

		channel, requests, err := ch.Accept()
		if err != nil {
			logger.Errorf("Could not accept channel: %v", err)
			continue
		}
		defer channel.Close()

		c.term = terminal.NewTerminal(channel, prompt)
		for req := range requests {
			var width, height int
			var ok bool

			switch req.Type {
			case "shell":
				if c.term != nil && !hasShell {
					go c.handleShell(channel)
					ok = true
					hasShell = true
				}
			case "pty-req":
				width, height, ok = parsePtyRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			case "window-change":
				width, height, ok = parseWinchRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			}

			if req.WantReply {
				req.Reply(ok, nil)
			}
		}
	}
}
<|endoftext|>