<commit_before>package tesla

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
)

type Auth struct {
	GrantType    string `json:"grant_type"`
	ClientID     string `json:"client_id"`
	ClientSecret string `json:"client_secret"`
	Email        string `json:"email"`
	Password     string `json:"password"`
	URL          string
	StreamingURL string
}

type Token struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
}

type Client struct {
	Auth  *Auth
	Token *Token
	HTTP  *http.Client
}

var (
	AuthURL      = "https://owner-api.teslamotors.com/oauth/token"
	BaseURL      = "https://owner-api.teslamotors.com/api/1"
	ActiveClient *Client
)

func NewClient(auth *Auth) (*Client, error) {
	if auth.URL == "" {
		auth.URL = BaseURL
	}
	if auth.StreamingURL == "" {
		auth.StreamingURL = StreamingURL
	}

	client := &Client{
		Auth: auth,
		HTTP: &http.Client{},
	}
	token, err := client.authorize(auth)
	if err != nil {
		return nil, err
	}
	client.Token = token
	ActiveClient = client
	return client, nil
}

func (c Client) authorize(auth *Auth) (*Token, error) {
	auth.GrantType = "password"
	data, _ := json.Marshal(auth)
	body, err := c.post(AuthURL, data)
	if err != nil {
		return nil, err
	}
	token := &Token{}
	err = json.Unmarshal(body, token)
	if err != nil {
		return nil, err
	}
	return token, nil
}

// // Calls an HTTP DELETE
func (c Client) delete(url string) error {
	req, _ := http.NewRequest("DELETE", url, nil)
	_, err := c.processRequest(req)
	return err
}

// Calls an HTTP GET
func (c Client) get(url string) ([]byte, error) {
	req, _ := http.NewRequest("GET", url, nil)
	return c.processRequest(req)
}

// Calls an HTTP POST with a JSON body
func (c Client) post(url string, body []byte) ([]byte, error) {
	req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body))
	return c.processRequest(req)
}

// // Calls an HTTP PUT
// func put(resource string, body []byte) ([]byte, error) {
// 	req, _ := http.NewRequest("PUT", BaseURL+resource, bytes.NewBuffer(body))
// 	return processRequest(req)
// }

// Processes a HTTP POST/PUT request
func (c Client) processRequest(req *http.Request) ([]byte, error) {
	c.setHeaders(req)
	res, err := c.HTTP.Do(req)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != 200 {
		return nil, errors.New(res.Status)
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	return body, nil
}

func (c Client) setHeaders(req *http.Request) {
	if c.Token != nil {
		req.Header.Set("Authorization", "Bearer "+c.Token.AccessToken)
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", "application/json")
}
<commit_msg>Added more comments for documenting the library as well as the put method.<commit_after>package tesla

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
)

// Required authorization credentials for the Tesla API
type Auth struct {
	GrantType    string `json:"grant_type"`
	ClientID     string `json:"client_id"`
	ClientSecret string `json:"client_secret"`
	Email        string `json:"email"`
	Password     string `json:"password"`
	URL          string
	StreamingURL string
}

// The token and related elements returned after a successful auth
// by the Tesla API
type Token struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
}

// Provides the client and associated elements for interacting with the
// Tesla API
type Client struct {
	Auth  *Auth
	Token *Token
	HTTP  *http.Client
}

var (
	AuthURL      = "https://owner-api.teslamotors.com/oauth/token"
	BaseURL      = "https://owner-api.teslamotors.com/api/1"
	ActiveClient *Client
)

// Generates a new client for the Tesla API
func NewClient(auth *Auth) (*Client, error) {
	if auth.URL == "" {
		auth.URL = BaseURL
	}
	if auth.StreamingURL == "" {
		auth.StreamingURL = StreamingURL
	}

	client := &Client{
		Auth: auth,
		HTTP: &http.Client{},
	}
	token, err := client.authorize(auth)
	if err != nil {
		return nil, err
	}
	client.Token = token
	ActiveClient = client
	return client, nil
}

// Authorizes against the Tesla API with the appropriate credentials
func (c Client) authorize(auth *Auth) (*Token, error) {
	auth.GrantType = "password"
	data, _ := json.Marshal(auth)
	body, err := c.post(AuthURL, data)
	if err != nil {
		return nil, err
	}
	token := &Token{}
	err = json.Unmarshal(body, token)
	if err != nil {
		return nil, err
	}
	return token, nil
}

// // Calls an HTTP DELETE
func (c Client) delete(url string) error {
	req, _ := http.NewRequest("DELETE", url, nil)
	_, err := c.processRequest(req)
	return err
}

// Calls an HTTP GET
func (c Client) get(url string) ([]byte, error) {
	req, _ := http.NewRequest("GET", url, nil)
	return c.processRequest(req)
}

// Calls an HTTP POST with a JSON body
func (c Client) post(url string, body []byte) ([]byte, error) {
	req, _ := http.NewRequest("POST", url, bytes.NewBuffer(body))
	return c.processRequest(req)
}

// Calls an HTTP PUT
func (c Client) put(resource string, body []byte) ([]byte, error) {
	req, _ := http.NewRequest("PUT", BaseURL+resource, bytes.NewBuffer(body))
	return c.processRequest(req)
}

// Processes a HTTP POST/PUT request
func (c Client) processRequest(req *http.Request) ([]byte, error) {
	c.setHeaders(req)
	res, err := c.HTTP.Do(req)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != 200 {
		return nil, errors.New(res.Status)
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	return body, nil
}

// Sets the required headers for calls to the Tesla API
func (c Client) setHeaders(req *http.Request) {
	if c.Token != nil {
		req.Header.Set("Authorization", "Bearer "+c.Token.AccessToken)
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", "application/json")
}
<|endoftext|>


<commit_before>// Package crowi provides some Crowi APIs for Go
package crowi

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"time"

	"github.com/pkg/errors"
)

const version = "0.1"

var userAgent = fmt.Sprintf("CrowiGoClient/%s (%s)", version, runtime.Version())

const (
	apiPagesCreate    = "/_api/pages.create"
	apiPagesUpdate    = "/_api/pages.update"
	apiAttachmentsAdd = "/_api/attachments.add"
)

type API interface {
	PagesCreate() (*Crowi, error)
	PagesUpdate() (*Crowi, error)
	AttachmentsAdd() (*Crowi, error)
}

// Client wraps http client
type Client struct {
	URL        *url.URL
	Token      string
	HTTPClient *http.Client
}

// NewClient creates an API client
func NewClient(apiURL, token string) (*Client, error) {
	if len(apiURL) == 0 {
		return nil, errors.New("missing api url")
	}

	if len(token) == 0 {
		return nil, errors.New("missing token")
	}

	parsedURL, err := url.ParseRequestURI(apiURL)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to parse url: %s", apiURL)
	}

	return &Client{
		URL:        parsedURL,
		Token:      token,
		HTTPClient: &http.Client{Timeout: 10 * time.Second},
	}, nil
}

func (c *Client) newRequest(method, resource string, data url.Values) (*http.Request, error) {
	c.URL.Path = resource
	urlStr := fmt.Sprintf("%v", c.URL)

	req, err := http.NewRequest(method, urlStr, bytes.NewBufferString(data.Encode()))
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("User-Agent", userAgent)
	req.Header.Set("Content-Length", strconv.Itoa(len(data.Encode())))

	return req, nil
}

// PagesCreate makes a page in your Crowi. The request requires
// the path and page content used for the page name
func (c *Client) PagesCreate(path, body string) (*Crowi, error) {
	data := url.Values{}
	data.Set("access_token", c.Token)
	data.Set("path", path)
	data.Set("body", body)

	req, err := c.newRequest("POST", apiPagesCreate, data)
	if err != nil {
		return nil, err
	}

	res, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}

	var crowi Crowi
	if err := decodeBody(res, &crowi); err != nil {
		return nil, err
	}

	return &crowi, nil
}

// PagesUpdate updates the page content. A page_id is necessary to know which
// page should be updated.
func (c *Client) PagesUpdate(pageID, body string) (*Crowi, error) {
	data := url.Values{}
	data.Set("access_token", c.Token)
	data.Set("page_id", pageID)
	data.Set("body", body)

	req, err := c.newRequest("POST", apiPagesUpdate, data)
	if err != nil {
		return nil, err
	}

	res, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}

	var crowi Crowi
	if err := decodeBody(res, &crowi); err != nil {
		return nil, err
	}

	return &crowi, nil
}

func (c *Client) fileUpload(method, resource string, params map[string]string, filePath string) (*http.Request, error) {
	c.URL.Path = resource
	urlStr := fmt.Sprintf("%v", c.URL)

	var buffer bytes.Buffer
	writer := multipart.NewWriter(&buffer)
	for key, val := range params {
		err := writer.WriteField(key, val)
		if err != nil {
			return nil, err
		}
	}
	{
		header := make(textproto.MIMEHeader)
		header.Add("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, filePath))
		header.Add("Content-Type", "image/png")
		fileWriter, err := writer.CreatePart(header)
		if err != nil {
			return nil, err
		}
		file, err := os.Open(filePath)
		if err != nil {
			return nil, err
		}
		io.Copy(fileWriter, file)
	}
	writer.Close()

	req, err := http.NewRequest(method, urlStr, &buffer)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "multipart/form-data; boundary="+writer.Boundary())
	req.Header.Set("User-Agent", userAgent)

	return req, nil
}

// AttachmentsAdd attaches an image file to the page. This request requires
// page_id and the image file path which you want to attach.
func (c *Client) AttachmentsAdd(pageID, filePath string) (*Crowi, error) {
	req, err := c.fileUpload("POST", apiAttachmentsAdd, map[string]string{
		"access_token": c.Token,
		"page_id":      pageID,
	}, filePath)
	if err != nil {
		return nil, err
	}

	res, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}

	var crowi Crowi
	if err := decodeBody(res, &crowi); err != nil {
		return nil, err
	}

	return &crowi, nil
}

func decodeBody(resp *http.Response, out interface{}) error {
	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	return decoder.Decode(out)
}
<commit_msg>Call String instead of fmt.Sprintf<commit_after>// Package crowi provides some Crowi APIs for Go
package crowi

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"time"

	"github.com/pkg/errors"
)

const version = "0.1"

var userAgent = fmt.Sprintf("CrowiGoClient/%s (%s)", version, runtime.Version())

const (
	apiPagesCreate    = "/_api/pages.create"
	apiPagesUpdate    = "/_api/pages.update"
	apiAttachmentsAdd = "/_api/attachments.add"
)

type API interface {
	PagesCreate() (*Crowi, error)
	PagesUpdate() (*Crowi, error)
	AttachmentsAdd() (*Crowi, error)
}

// Client wraps http client
type Client struct {
	URL        *url.URL
	Token      string
	HTTPClient *http.Client
}

// NewClient creates an API client
func NewClient(apiURL, token string) (*Client, error) {
	if len(apiURL) == 0 {
		return nil, errors.New("missing api url")
	}

	if len(token) == 0 {
		return nil, errors.New("missing token")
	}

	parsedURL, err := url.ParseRequestURI(apiURL)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to parse url: %s", apiURL)
	}

	return &Client{
		URL:        parsedURL,
		Token:      token,
		HTTPClient: &http.Client{Timeout: 10 * time.Second},
	}, nil
}

func (c *Client) newRequest(method, resource string, data url.Values) (*http.Request, error) {
	c.URL.Path = resource
	urlStr := c.URL.String()

	req, err := http.NewRequest(method, urlStr, bytes.NewBufferString(data.Encode()))
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("User-Agent", userAgent)
	req.Header.Set("Content-Length", strconv.Itoa(len(data.Encode())))

	return req, nil
}

// PagesCreate makes a page in your Crowi. The request requires
// the path and page content used for the page name
func (c *Client) PagesCreate(path, body string) (*Crowi, error) {
	data := url.Values{}
	data.Set("access_token", c.Token)
	data.Set("path", path)
	data.Set("body", body)

	req, err := c.newRequest("POST", apiPagesCreate, data)
	if err != nil {
		return nil, err
	}

	res, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}

	var crowi Crowi
	if err := decodeBody(res, &crowi); err != nil {
		return nil, err
	}

	return &crowi, nil
}

// PagesUpdate updates the page content. A page_id is necessary to know which
// page should be updated.
func (c *Client) PagesUpdate(pageID, body string) (*Crowi, error) {
	data := url.Values{}
	data.Set("access_token", c.Token)
	data.Set("page_id", pageID)
	data.Set("body", body)

	req, err := c.newRequest("POST", apiPagesUpdate, data)
	if err != nil {
		return nil, err
	}

	res, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}

	var crowi Crowi
	if err := decodeBody(res, &crowi); err != nil {
		return nil, err
	}

	return &crowi, nil
}

func (c *Client) fileUpload(method, resource string, params map[string]string, filePath string) (*http.Request, error) {
	c.URL.Path = resource
	urlStr := c.URL.String()

	var buffer bytes.Buffer
	writer := multipart.NewWriter(&buffer)
	for key, val := range params {
		err := writer.WriteField(key, val)
		if err != nil {
			return nil, err
		}
	}
	{
		header := make(textproto.MIMEHeader)
		header.Add("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, filePath))
		header.Add("Content-Type", "image/png")
		fileWriter, err := writer.CreatePart(header)
		if err != nil {
			return nil, err
		}
		file, err := os.Open(filePath)
		if err != nil {
			return nil, err
		}
		io.Copy(fileWriter, file)
	}
	writer.Close()

	req, err := http.NewRequest(method, urlStr, &buffer)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "multipart/form-data; boundary="+writer.Boundary())
	req.Header.Set("User-Agent", userAgent)

	return req, nil
}

// AttachmentsAdd attaches an image file to the page. This request requires
// page_id and the image file path which you want to attach.
func (c *Client) AttachmentsAdd(pageID, filePath string) (*Crowi, error) {
	req, err := c.fileUpload("POST", apiAttachmentsAdd, map[string]string{
		"access_token": c.Token,
		"page_id":      pageID,
	}, filePath)
	if err != nil {
		return nil, err
	}

	res, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, err
	}

	var crowi Crowi
	if err := decodeBody(res, &crowi); err != nil {
		return nil, err
	}

	return &crowi, nil
}

func decodeBody(resp *http.Response, out interface{}) error {
	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	return decoder.Decode(out)
}
<|endoftext|>


<commit_before>package livestatus

import (
	"net"
)

const bufferSize = 1024

// Client represents a Livestatus client instance.
type Client struct {
	network string
	address string
	dialer  *net.Dialer
	conn    net.Conn
}

// NewClient creates a new Livestatus client instance.
func NewClient(network, address string) *Client {
	return NewClientWithDialer(network, address, new(net.Dialer))
}

// NewClientWithDialer creates a new Livestatus client instance using a provided network dialer.
func NewClientWithDialer(network, address string, dialer *net.Dialer) *Client {
	return &Client{
		network: network,
		address: address,
		dialer:  dialer,
	}
}

// Close closes any remaining connection.
func (c *Client) Close() {
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}
}

// Exec executes a given Livestatus query.
func (c *Client) Exec(r Request) (*Response, error) {
	var err error

	// Initialize connection if none available
	if c.conn == nil {
		c.conn, err = c.dialer.Dial(c.network, c.address)
		if err != nil {
			return nil, err
		}

		if !r.keepAlive() {
			defer c.Close()
		}
	}

	return r.handle(c.conn)
}
<commit_msg>Add missing keepalive on TCP connection<commit_after>package livestatus

import (
	"net"
)

const bufferSize = 1024

// Client represents a Livestatus client instance.
type Client struct {
	network string
	address string
	dialer  *net.Dialer
	conn    net.Conn
}

// NewClient creates a new Livestatus client instance.
func NewClient(network, address string) *Client {
	return NewClientWithDialer(network, address, new(net.Dialer))
}

// NewClientWithDialer creates a new Livestatus client instance using a provided network dialer.
func NewClientWithDialer(network, address string, dialer *net.Dialer) *Client {
	return &Client{
		network: network,
		address: address,
		dialer:  dialer,
	}
}

// Close closes any remaining connection.
func (c *Client) Close() {
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}
}

// Exec executes a given Livestatus query.
func (c *Client) Exec(r Request) (*Response, error) {
	var err error

	// Initialize connection if none available
	if c.conn == nil {
		c.conn, err = c.dialer.Dial(c.network, c.address)
		if err != nil {
			return nil, err
		}

		if r.keepAlive() {
			c.conn.(*net.TCPConn).SetKeepAlive(true)
		} else {
			defer c.Close()
		}
	}

	return r.handle(c.conn)
}
<|endoftext|>


<commit_before>package ratsit

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
)

const (
	// APIURL defines the Ratsit API URL for use in production
	APIURL           = "https://api.ratsit.se/api/v1"
	pkgPersonSearch  = "personsok"
	pkgCompanySearch = "foretagsok"
)

var (
	// ErrInvalidInput ...
	ErrInvalidInput = errors.New("invalid input")
	// ErrInternalServer ...
	ErrInternalServer = errors.New("internal server error")
	// ErrInvalidCredentials ...
	ErrInvalidCredentials = errors.New("authentication failed")
)

// Ratsit is the the client
type Ratsit struct {
	apiURL string
	apiKey string
	client *http.Client
}

// New creates a new client to interact with the Ratsit API
func New(apiURL string, key string) (r Ratsit) {
	r.apiURL = apiURL
	r.apiKey = key
	r.client = new(http.Client)
	return
}

// GetPerson returns a person from the database by looking up their unique personnummer
func (r *Ratsit) GetPerson(ssn string, pkg string) (person Person, err error) {
	// TODO: Validate SSN and pkg
	url := generatePersonInformationURL(r.apiURL, ssn)
	body, err := doHTTPRequest(r.client, url, r.apiKey, pkg)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, &person)
	return
}

// SearchPerson searches the Ratsit database for people with the name and location given in the parameters
func (r *Ratsit) SearchPerson(name string, location string, limit int, recordFrom int) (personSearchResults SearchResults, err error) {
	url := generatePersonSearchURL(r.apiURL, name, location, limit, recordFrom)
	body, err := doHTTPRequest(r.client, url, r.apiKey, pkgPersonSearch)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, &personSearchResults)
	return
}

// GetCompany returns a company from the database by looking up their unique organization number
func (r *Ratsit) GetCompany(organizationNumber string, pkg string) (company CompanyInformationResponse, err error) {
	// TODO: Validate organizationNumber and pkg
	url := generateCompanyInformationURL(r.apiURL, organizationNumber)
	body, err := doHTTPRequest(r.client, url, r.apiKey, pkg)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, &company)
	return
}

// SearchCompany searches the Ratsit database for companies with the name and location given in the parameters
func (r *Ratsit) SearchCompany(name string, location string, limit int, recordFrom int) (companySearchResults CompanySearchResults, err error) {
	url := generateCompanySearchURL(r.apiURL, name, location, limit, recordFrom)
	body, err := doHTTPRequest(r.client, url, r.apiKey, pkgCompanySearch)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, &companySearchResults)
	return
}

func doHTTPRequest(client *http.Client, url string, apiKey string, pkg string) (body []byte, err error) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return
	}
	authorizeRequest(req, apiKey, pkg)
	resp, err := client.Do(req)
	if err != nil {
		return
	}
	err = handleResponseError(resp)
	if err != nil {
		return
	}
	body, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return
	}
	return
}
<commit_msg>update ratsit api endpoint<commit_after>package ratsit

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
)

const (
	// APIURL defines the Ratsit API URL for use in production
	APIURL           = "https://api.checkbiz.se/api/v1"
	pkgPersonSearch  = "personsok"
	pkgCompanySearch = "foretagsok"
)

var (
	// ErrInvalidInput ...
	ErrInvalidInput = errors.New("invalid input")
	// ErrInternalServer ...
	ErrInternalServer = errors.New("internal server error")
	// ErrInvalidCredentials ...
	ErrInvalidCredentials = errors.New("authentication failed")
)

// Ratsit is the the client
type Ratsit struct {
	apiURL string
	apiKey string
	client *http.Client
}

// New creates a new client to interact with the Ratsit API
func New(apiURL string, key string) (r Ratsit) {
	r.apiURL = apiURL
	r.apiKey = key
	r.client = new(http.Client)
	return
}

// GetPerson returns a person from the database by looking up their unique personnummer
func (r *Ratsit) GetPerson(ssn string, pkg string) (person Person, err error) {
	// TODO: Validate SSN and pkg
	url := generatePersonInformationURL(r.apiURL, ssn)
	body, err := doHTTPRequest(r.client, url, r.apiKey, pkg)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, &person)
	return
}

// SearchPerson searches the Ratsit database for people with the name and location given in the parameters
func (r *Ratsit) SearchPerson(name string, location string, limit int, recordFrom int) (personSearchResults SearchResults, err error) {
	url := generatePersonSearchURL(r.apiURL, name, location, limit, recordFrom)
	body, err := doHTTPRequest(r.client, url, r.apiKey, pkgPersonSearch)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, &personSearchResults)
	return
}

// GetCompany returns a company from the database by looking up their unique organization number
func (r *Ratsit) GetCompany(organizationNumber string, pkg string) (company CompanyInformationResponse, err error) {
	// TODO: Validate organizationNumber and pkg
	url := generateCompanyInformationURL(r.apiURL, organizationNumber)
	body, err := doHTTPRequest(r.client, url, r.apiKey, pkg)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, &company)
	return
}

// SearchCompany searches the Ratsit database for companies with the name and location given in the parameters
func (r *Ratsit) SearchCompany(name string, location string, limit int, recordFrom int) (companySearchResults CompanySearchResults, err error) {
	url := generateCompanySearchURL(r.apiURL, name, location, limit, recordFrom)
	body, err := doHTTPRequest(r.client, url, r.apiKey, pkgCompanySearch)
	if err != nil {
		return
	}
	err = json.Unmarshal(body, &companySearchResults)
	return
}

func doHTTPRequest(client *http.Client, url string, apiKey string, pkg string) (body []byte, err error) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return
	}
	authorizeRequest(req, apiKey, pkg)
	resp, err := client.Do(req)
	if err != nil {
		return
	}
	err = handleResponseError(resp)
	if err != nil {
		return
	}
	body, err = ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/iplist"
	"github.com/dustin/go-humanize"
)

const clearScreen = "\033[H\033[2J"

const torrentBlockListURL = "http://john.bitsurge.net/public/biglist.p2p.gz"

var isHTTP = regexp.MustCompile(`^https?:\/\/`)

// ClientError formats errors coming from the client.
type ClientError struct {
	Type   string
	Origin error
}

func (clientError ClientError) Error() string {
	return fmt.Sprintf("Error %s: %s\n", clientError.Type, clientError.Origin)
}

// Client manages the torrent downloading.
type Client struct {
	Client   *torrent.Client
	Torrent  torrent.Torrent
	Progress int64
	Port     int
}

// NewClient creates a new torrent client based on a magnet or a torrent file.
// If the torrent file is on http, we try downloading it.
func NewClient(torrentPath string, port int, seed bool, tcp bool) (client Client, err error) {
	var t torrent.Torrent
	var c *torrent.Client

	client.Port = port

	// Create client.
	c, err = torrent.NewClient(&torrent.Config{
		DataDir:    os.TempDir(),
		NoUpload:   !seed,
		Seed:       seed,
		DisableTCP: !tcp,
	})

	if err != nil {
		return client, ClientError{Type: "creating torrent client", Origin: err}
	}

	client.Client = c

	// Add torrent.

	// Add as magnet url.
	if strings.HasPrefix(torrentPath, "magnet:") {
		if t, err = c.AddMagnet(torrentPath); err != nil {
			return client, ClientError{Type: "adding torrent", Origin: err}
		}
	} else {
		// Otherwise add as a torrent file.

		// If it's online, we try downloading the file.
		if isHTTP.MatchString(torrentPath) {
			if torrentPath, err = downloadFile(torrentPath); err != nil {
				return client, ClientError{Type: "downloading torrent file", Origin: err}
			}
		}

		// Check if the file exists.
		if _, err = os.Stat(torrentPath); err != nil {
			return client, ClientError{Type: "file not found", Origin: err}
		}

		if t, err = c.AddTorrentFromFile(torrentPath); err != nil {
			return client, ClientError{Type: "adding torrent to the client", Origin: err}
		}
	}

	client.Torrent = t

	go func() {
		<-t.GotInfo()
		t.DownloadAll()

		// Prioritize first 5% of the file.
		client.getLargestFile().PrioritizeRegion(0, int64(t.NumPieces()/100*5))
	}()

	go client.addBlocklist()

	return
}

// Download and add the blocklist.
func (c *Client) addBlocklist() {
	if c.Client.IPBlockList() != nil {
		log.Printf("Found blocklist")
		return
	}

	var err error
	blocklistPath := c.Client.ConfigDir() + "/blocklist"

	// Download blocklist.
	log.Printf("Downloading blocklist")
	fileName, err := downloadFile(torrentBlockListURL)
	if err != nil {
		log.Printf("Error downloading blocklist: %s\n", err)
		return
	}

	// Ungzip file.
	in, err := os.Open(fileName)
	if err != nil {
		log.Printf("Error extracting blocklist: %s\n", err)
		return
	}
	defer func() {
		if err = in.Close(); err != nil {
			log.Printf("Error closing the blocklist gzip file: %s", err)
		}
	}()
	reader, err := gzip.NewReader(in)

	// Write to {configdir}/blocklist
	out, err := os.Create(blocklistPath)
	if err != nil {
		log.Printf("Error writing blocklist: %s\n", err)
		return
	}
	defer func() {
		if err = out.Close(); err != nil {
			log.Printf("Error writing the blocklist file: %s", err)
		}
	}()
	_, err = io.Copy(out, reader)
	if err != nil {
		log.Printf("Error writing the blocklist file: %s", err)
		return
	}

	// Add to the blocklist.
	blocklistReader, err := os.Open(blocklistPath)
	blocklist, err := iplist.NewFromReader(blocklistReader)
	if err != nil {
		log.Printf("Error reading blocklist: %s", err)
		return
	}
	log.Printf("Setting blocklist.\nFound %d ranges\n", blocklist.NumRanges())
	c.Client.SetIPBlockList(blocklist)
}

// Close cleans up the connections.
func (c *Client) Close() {
	c.Torrent.Drop()
	c.Client.Close()
}

// Render outputs the command line interface for the client.
func (c *Client) Render() {
	t := c.Torrent

	if t.Info() == nil {
		return
	}

	var currentProgress = t.BytesCompleted()
	speed := humanize.Bytes(uint64(currentProgress-c.Progress)) + "/s"
	c.Progress = currentProgress

	complete := humanize.Bytes(uint64(currentProgress))
	size := humanize.Bytes(uint64(t.Info().TotalLength()))

	print(clearScreen)
	fmt.Println(t.Info().Name)
	fmt.Println("=============================================================")
	if c.ReadyForPlayback() {
		fmt.Printf("Stream: \thttp://localhost:%d\n", c.Port)
	}

	if currentProgress > 0 {
		fmt.Printf("Progress: \t%s / %s  %.2f%%\n", complete, size, c.percentage())
	}
	if currentProgress < t.Info().TotalLength() {
		fmt.Printf("Download speed: %s\n", speed)
	}
	//fmt.Printf("Connections: \t%d\n", len(t.Conns))
	//fmt.Printf("%s\n", c.RenderPieces())
}

func (c Client) getLargestFile() *torrent.File {
	var target torrent.File
	var maxSize int64

	for _, file := range c.Torrent.Files() {
		if maxSize < file.Length() {
			maxSize = file.Length()
			target = file
		}
	}

	return &target
}

/*
func (c Client) RenderPieces() (output string) {
	pieces := c.Torrent.PieceStateRuns()
	for i := range pieces {
		piece := pieces[i]

		if piece.Priority == torrent.PiecePriorityReadahead {
			output += "!"
		}

		if piece.Partial {
			output += "P"
		} else if piece.Checking {
			output += "c"
		} else if piece.Complete {
			output += "d"
		} else {
			output += "_"
		}
	}

	return
}
*/

// ReadyForPlayback checks if the torrent is ready for playback or not.
// We wait until 5% of the torrent to start playing.
func (c Client) ReadyForPlayback() bool {
	return c.percentage() > 5
}

// GetFile is an http handler to serve the biggest file managed by the client.
func (c Client) GetFile(w http.ResponseWriter, r *http.Request) {
	target := c.getLargestFile()
	entry, err := NewFileReader(target)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	defer func() {
		if err := entry.Close(); err != nil {
			log.Printf("Error closing file reader: %s\n", err)
		}
	}()

	w.Header().Set("Content-Disposition", "attachment; filename=\""+c.Torrent.Info().Name+"\"")
	http.ServeContent(w, r, target.DisplayPath(), time.Now(), entry)
}

func (c Client) percentage() float64 {
	info := c.Torrent.Info()

	if info == nil {
		return 0
	}

	return float64(c.Torrent.BytesCompleted()) / float64(info.TotalLength()) * 100
}

func downloadFile(URL string) (fileName string, err error) {
	var file *os.File
	if file, err = ioutil.TempFile(os.TempDir(), "torrent-imageviewer"); err != nil {
		return
	}

	defer func() {
		if err = file.Close(); err != nil {
			log.Printf("Error closing torrent file: %s", err)
		}
	}()

	response, err := http.Get(URL)
	if err != nil {
		return
	}

	defer func() {
		if err = response.Body.Close(); err != nil {
			log.Printf("Error closing torrent file: %s", err)
		}
	}()

	_, err = io.Copy(file, response.Body)

	return file.Name(), err
}
<commit_msg>Rename temporary torrent file variable<commit_after>package main

import (
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/iplist"
	"github.com/dustin/go-humanize"
)

const clearScreen = "\033[H\033[2J"

const torrentBlockListURL = "http://john.bitsurge.net/public/biglist.p2p.gz"

var isHTTP = regexp.MustCompile(`^https?:\/\/`)

// ClientError formats errors coming from the client.
type ClientError struct {
	Type   string
	Origin error
}

func (clientError ClientError) Error() string {
	return fmt.Sprintf("Error %s: %s\n", clientError.Type, clientError.Origin)
}

// Client manages the torrent downloading.
type Client struct {
	Client   *torrent.Client
	Torrent  torrent.Torrent
	Progress int64
	Port     int
}

// NewClient creates a new torrent client based on a magnet or a torrent file.
// If the torrent file is on http, we try downloading it.
func NewClient(torrentPath string, port int, seed bool, tcp bool) (client Client, err error) {
	var t torrent.Torrent
	var c *torrent.Client

	client.Port = port

	// Create client.
	c, err = torrent.NewClient(&torrent.Config{
		DataDir:    os.TempDir(),
		NoUpload:   !seed,
		Seed:       seed,
		DisableTCP: !tcp,
	})

	if err != nil {
		return client, ClientError{Type: "creating torrent client", Origin: err}
	}

	client.Client = c

	// Add torrent.

	// Add as magnet url.
	if strings.HasPrefix(torrentPath, "magnet:") {
		if t, err = c.AddMagnet(torrentPath); err != nil {
			return client, ClientError{Type: "adding torrent", Origin: err}
		}
	} else {
		// Otherwise add as a torrent file.

		// If it's online, we try downloading the file.
		if isHTTP.MatchString(torrentPath) {
			if torrentPath, err = downloadFile(torrentPath); err != nil {
				return client, ClientError{Type: "downloading torrent file", Origin: err}
			}
		}

		// Check if the file exists.
		if _, err = os.Stat(torrentPath); err != nil {
			return client, ClientError{Type: "file not found", Origin: err}
		}

		if t, err = c.AddTorrentFromFile(torrentPath); err != nil {
			return client, ClientError{Type: "adding torrent to the client", Origin: err}
		}
	}

	client.Torrent = t

	go func() {
		<-t.GotInfo()
		t.DownloadAll()

		// Prioritize first 5% of the file.
		client.getLargestFile().PrioritizeRegion(0, int64(t.NumPieces()/100*5))
	}()

	go client.addBlocklist()

	return
}

// Download and add the blocklist.
func (c *Client) addBlocklist() {
	if c.Client.IPBlockList() != nil {
		log.Printf("Found blocklist")
		return
	}

	var err error
	blocklistPath := c.Client.ConfigDir() + "/blocklist"

	// Download blocklist.
	log.Printf("Downloading blocklist")
	fileName, err := downloadFile(torrentBlockListURL)
	if err != nil {
		log.Printf("Error downloading blocklist: %s\n", err)
		return
	}

	// Ungzip file.
	in, err := os.Open(fileName)
	if err != nil {
		log.Printf("Error extracting blocklist: %s\n", err)
		return
	}
	defer func() {
		if err = in.Close(); err != nil {
			log.Printf("Error closing the blocklist gzip file: %s", err)
		}
	}()
	reader, err := gzip.NewReader(in)

	// Write to {configdir}/blocklist
	out, err := os.Create(blocklistPath)
	if err != nil {
		log.Printf("Error writing blocklist: %s\n", err)
		return
	}
	defer func() {
		if err = out.Close(); err != nil {
			log.Printf("Error writing the blocklist file: %s", err)
		}
	}()
	_, err = io.Copy(out, reader)
	if err != nil {
		log.Printf("Error writing the blocklist file: %s", err)
		return
	}

	// Add to the blocklist.
	blocklistReader, err := os.Open(blocklistPath)
	blocklist, err := iplist.NewFromReader(blocklistReader)
	if err != nil {
		log.Printf("Error reading blocklist: %s", err)
		return
	}
	log.Printf("Setting blocklist.\nFound %d ranges\n", blocklist.NumRanges())
	c.Client.SetIPBlockList(blocklist)
}

// Close cleans up the connections.
func (c *Client) Close() {
	c.Torrent.Drop()
	c.Client.Close()
}

// Render outputs the command line interface for the client.
func (c *Client) Render() {
	t := c.Torrent

	if t.Info() == nil {
		return
	}

	var currentProgress = t.BytesCompleted()
	speed := humanize.Bytes(uint64(currentProgress-c.Progress)) + "/s"
	c.Progress = currentProgress

	complete := humanize.Bytes(uint64(currentProgress))
	size := humanize.Bytes(uint64(t.Info().TotalLength()))

	print(clearScreen)
	fmt.Println(t.Info().Name)
	fmt.Println("=============================================================")
	if c.ReadyForPlayback() {
		fmt.Printf("Stream: \thttp://localhost:%d\n", c.Port)
	}

	if currentProgress > 0 {
		fmt.Printf("Progress: \t%s / %s  %.2f%%\n", complete, size, c.percentage())
	}
	if currentProgress < t.Info().TotalLength() {
		fmt.Printf("Download speed: %s\n", speed)
	}
	//fmt.Printf("Connections: \t%d\n", len(t.Conns))
	//fmt.Printf("%s\n", c.RenderPieces())
}

func (c Client) getLargestFile() *torrent.File {
	var target torrent.File
	var maxSize int64

	for _, file := range c.Torrent.Files() {
		if maxSize < file.Length() {
			maxSize = file.Length()
			target = file
		}
	}

	return &target
}

/*
func (c Client) RenderPieces() (output string) {
	pieces := c.Torrent.PieceStateRuns()
	for i := range pieces {
		piece := pieces[i]

		if piece.Priority == torrent.PiecePriorityReadahead {
			output += "!"
		}

		if piece.Partial {
			output += "P"
		} else if piece.Checking {
			output += "c"
		} else if piece.Complete {
			output += "d"
		} else {
			output += "_"
		}
	}

	return
}
*/

// ReadyForPlayback checks if the torrent is ready for playback or not.
// We wait until 5% of the torrent to start playing.
func (c Client) ReadyForPlayback() bool {
	return c.percentage() > 5
}

// GetFile is an http handler to serve the biggest file managed by the client.
func (c Client) GetFile(w http.ResponseWriter, r *http.Request) {
	target := c.getLargestFile()
	entry, err := NewFileReader(target)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	defer func() {
		if err := entry.Close(); err != nil {
			log.Printf("Error closing file reader: %s\n", err)
		}
	}()

	w.Header().Set("Content-Disposition", "attachment; filename=\""+c.Torrent.Info().Name+"\"")
	http.ServeContent(w, r, target.DisplayPath(), time.Now(), entry)
}

func (c Client) percentage() float64 {
	info := c.Torrent.Info()

	if info == nil {
		return 0
	}

	return float64(c.Torrent.BytesCompleted()) / float64(info.TotalLength()) * 100
}

func downloadFile(URL string) (fileName string, err error) {
	var file *os.File
	if file, err = ioutil.TempFile(os.TempDir(), "go-peerflix"); err != nil {
		return
	}

	defer func() {
		if err = file.Close(); err != nil {
			log.Printf("Error closing torrent file: %s", err)
		}
	}()

	response, err := http.Get(URL)
	if err != nil {
		return
	}

	defer func() {
		if err = response.Body.Close(); err != nil {
			log.Printf("Error closing torrent file: %s", err)
		}
	}()

	_, err = io.Copy(file, response.Body)

	return file.Name(), err
}
<|endoftext|>


<commit_before>package mqttclient

import (
	"crypto/tls"
	"errors"
	"fmt"
	mqtt "github.com/clearblade/mqtt_parsing"
	"io"
	mrand "math/rand"
	"net"
	"sync"
	"time"
)

type Client struct {

	//this is all clearblade specific
	SystemKey    string
	SystemSecret string
	AuthToken    string
	Clientid     string
	//this will usually be occupied
	//by a net.Conn
	C                   io.ReadWriteCloser
	internalOutgoingBuf chan []byte
	Timeout             time.Duration
	//thou shalt type channels consumed by others
	ClientErrorBuffer chan error

	internalErrorBuffer chan *errWrap
	shutdown_reader     chan struct{}
	shutdown_writer     chan struct{}
	//introduce a sync write mode?
	last_timeout_reccd time.Time

	shutting_down bool

	msg_store                *storage
	subscriptions            *outgoing_topics
	waiting_for_subscription *subscription_store
	//TODO:redesign around a thread-local
	//rng
	rando    *mrand.Rand
	randomut *sync.RWMutex
}

var (
	Verbose bool
)

//Start connects to the mqtt broker. It does not send the connect packet. Use the SendConnect function for that.
func (c *Client) Start(addr string, ssl *tls.Config) error {
	var con net.Conn
	var err error
	if ssl != nil {
		con, err = tls.Dial("tcp", addr, ssl)
	} else {
		con, err = net.Dial("tcp", addr)
	}
	if err != nil {
		return err
	}
	c.C = con

	go c.connectionWriter()
	go c.connectionListener()
	go c.errorTree()
	return nil
}

//NewClient allocates a new client. It is supplied with the (in order of appearance)
//Token, SystemKey,SystemSecret,Clientid, and the mqtt timeout
//Note that the following combinations of (Token|SystemKey|SystemSecret) are allowed
//(Token && SystemKey), (SystemKey && SystemSecret)
func NewClient(tok, sk, ss, cid string, timeout int) *Client {
	client := &Client{
		msg_store:                newStorage(),
		waiting_for_subscription: newSubscriptionStore(),
		subscriptions:            newOutgoingTopics(),
		SystemSecret:             ss,
		SystemKey:                sk,
		AuthToken:                tok,
		Clientid:                 cid,
		Timeout:                  time.Duration(timeout) * time.Second,
		internalOutgoingBuf:      make(chan []byte, 30),
		ClientErrorBuffer:        make(chan error, 10),
		internalErrorBuffer:      make(chan *errWrap, 2),
		shutdown_reader:          make(chan struct{}, 1),
		shutdown_writer:          make(chan struct{}, 1),
		rando:                    mrand.New(mrand.NewSource(time.Now().UnixNano())),
		randomut:                 new(sync.RWMutex),
	}
	return client
}

//sendMessage is an internal function that acts as a central point
//of failure for all of the message sending channels
//sort of like a fan-in, except this simply allows us to do
// all of the error handling logic in one place
func (c *Client) sendMessage(m mqtt.Message) error {
	select {
	case c.internalOutgoingBuf <- m.Encode():
		return nil
	default:
		return errors.New(fmt.Sprintf("Sending Channel blocked at %d", len(c.internalOutgoingBuf)))
	}
}

//connectionWriter is an internal function. it essentially sits in a goroutine and writes to the connection
//whenever it recieves a message over the channel
func (c *Client) connectionWriter() {
	if c.C == nil {
		return
	}
	for {
		select {
		case out := <-c.internalOutgoingBuf:
			_, err := c.C.Write(out)
			if err != nil {
				c.internalErrorBuffer <- &errWrap{
					err:      err,
					reciever: _CON_WRITER,
				}
				return
			}
		case <-c.shutdown_writer:
			return
		}
	}
}

//connectionListener is another function that sits on a goroutine.
//DecodePacket blocks until it reads a complete mqtt packet
func (c *Client) connectionListener() {
	if c.C == nil {
		return
	}
	mch, ech, shutdown := make(chan mqtt.Message, 10), make(chan error, 1), false
	//we have to establish an internal chain of goroutines here
	//otherwise we couldn't shutdown the listener on demand
	//since it's really hard to coordinate all the shutting down
	//when a connection drops
	//we're waiting for the connection listener to simply fail
	//this allows us to handle it a bit more gracefully
	//in order to shut down via channels directly we'd have to
	//wait for the read to fail anyway.
	go func(m chan mqtt.Message, e chan error) {
		for {
			msg, err := mqtt.DecodePacket(c.C)
			if err != nil {
				e <- err
				return
			}
			m <- msg
		}

	}(mch, ech)

	for {
		select {

		case msg := <-mch:
			//dispatch the internal functions
			c.dispatch(msg)
		case e := <-ech:
			//an error was recieved from the listenr
			if !shutdown {
				c.internalErrorBuffer <- &errWrap{
					err:      e,
					reciever: _CON_READER,
				}
			}
		case <-time.After(c.Timeout):
			//keepalive
			if time.Since(c.last_timeout_reccd) > c.Timeout {
				//there was a timeout
				c.Shutdown(true)
			}
			c.sendMessage(&mqtt.Pingreq{})

		case <-c.shutdown_reader:
			shutdown = true
			return
		}
	}
}

//this is the "business logic" of the client. it decides how each packet mutates the
//internal state of the client
func (c *Client) dispatch(msg mqtt.Message) {
	//for example we make a decision if the client sees this request or not
	//or if we have to send another message in a flow
	c.last_timeout_reccd = time.Now()
	switch msg.Type() {
	case mqtt.CONNECT:
		//shouldn't happen?
	case mqtt.CONNACK:
	case mqtt.PUBLISH:
		c.subscriptions.relay_message(msg.(*mqtt.Publish), msg.(*mqtt.Publish).Topic.Whole)
		switch msg.(*mqtt.Publish).Header.QOS {
		case 1:
			c.sendMessage(&mqtt.Puback{
				MessageId: msg.(*mqtt.Publish).MessageId,
			})
		case 2:
			c.sendMessage(&mqtt.Pubrec{
				MessageId: msg.(*mqtt.Publish).MessageId,
			})
		}
	case mqtt.PUBACK:
		//TODO:handle resend
	case mqtt.PUBREC:
		//discard, store the fact that it was recieved
		c.sendMessage(&mqtt.Pubrel{
			MessageId: msg.(*mqtt.Pubrec).MessageId,
			Header: &mqtt.StaticHeader{
				DUP:    false,
				Retain: false,
				QOS:    1,
			}})
	case mqtt.PUBREL:
		//this shouldn't have happened
	case mqtt.SUBSCRIBE:
		//this is not supposed to happen
	case mqtt.SUBACK:
		//the subscribe call blocks, so we need to forward the message
		//along that the subscribe was acknowleged so we can return
		//control flow to the parent program
		//of course, the problem is that a suback does not have
		//the subscriptions in it by name
		//but it does have the same message id
		//so we have to retrieve that and then match up the subscribe
		//TODO:NOTE THAT WE ARE ONLY USING ONE TOPIC PER SUBSCRIBE MESSSAGE
		//THIS LOGIC WILL NEED TWEAKING IF THAT CHANGES
		msg := c.msg_store.getEntry(msg.(*mqtt.Suback).MessageId)
		if msg == nil {
			//this is a bad thing to happen
			return
		}
		sub, ok := msg.(*mqtt.Subscribe)
		if !ok {
			//this is a worse thing to happen
			return
		}
		//TODO: BUG:: IF WE MAKE MULTISUBSCRIPTION, THIS WILL BREAK
		//we've now released control flow in that channel
		//also it'll allocate the userside channel and all that good stuff
		c.waiting_for_subscription.relay_message(msg, sub.Subscriptions[0].Topic.Whole)

	case mqtt.UNSUBSCRIBE:
		//this shouldn't happen
	case mqtt.UNSUBACK:
		//not gorgeous, but it do the thing
		msg := c.msg_store.getEntry(msg.(*mqtt.Unsuback).MessageId)
		unsub, ok := msg.(*mqtt.Unsubscribe)
		if !ok {
			return
		}
		c.subscriptions.remove_subscription(unsub.Topics[0].Topic.Whole)
		//TODO: do something besides prepend UNSUBSCRIBE# (the hash makes it an invalid mqtt topic, which should prevent collisions) to the front of the topic
		//but also doesn't allocate another big-ole map wrapper
		err := c.waiting_for_subscription.relay_message(msg, "UNSUBSCRIBE#"+unsub.Topics[0].Topic.Whole)
		if err != nil {
			c.ClientErrorBuffer <- err
		}
	case mqtt.PINGREQ:
		//shouldn't happen
	case mqtt.PINGRESP:
		//pingresp will reset the counter elsewhere
	case mqtt.DISCONNECT:
		//shouldn't happen
	default:
		c.ClientErrorBuffer <- fmt.Errorf("Invalid mqtt type recieved %+v", msg)
	}
}

//errorTree is the goroutine that sits on it's own goroutine and waits for
//a message to be recieved on c.internalErrorBuffer. It's our "in case of emergency break glass"
//way of reporting an error, and shutting the entire thing down
func (c *Client) errorTree() {
	//we still need to shutdown the listeners anyway
	//at least write will probably not error out

	//since you're reading the text of this fn, prepare your face for some exposition on how this mechanism works
	//so, we've spread reading and writing to the conn (or whatever) across goroutines, this is great
	//high speed low drag
	//but what happens if one goroutine encounters an error? the goroutines don't know about each other, so what do we do?
	//well, writing, and reading from to a closed connection is an error condition. so if one goroutine dies, then the other
	//will be taken down with it
	//we also use this mechanism for a regular shutdown of the client's connection, simply crashing them both and releasing the resources
	e := <-c.internalErrorBuffer
	if e.reciever != _CON_READER {
		c.shutdown_reader <- struct{}{}
	}
	if e.reciever != _CON_WRITER {
		c.shutdown_writer <- struct{}{}
	}
	if c.C != nil {
		c.C.Close()
	}
}

//Shutdown sends a disconnect packet (if asked), and then disconnects from the broker after a set time limit
func (c *Client) Shutdown(sendDisconnect bool) error {
	var err error
	if sendDisconnect {
		e := SendDisconnect(c)
		if e != nil {
			//don't return here, wait to finish the flow
			err = errors.New("While sending disconnect: " + e.Error() + "\nNote:connection was shut down anyway")
		}
		<-time.After(time.Second)
	}
	c.internalErrorBuffer <- &errWrap{reciever: _REGULAR_SHUTDOWN}
	return err
}

func (c *Client) randoPerm(i int) []int {
	c.randomut.RLock()
	order := c.rando.Perm(i)
	c.randomut.RUnlock()
	return order
}

func (c *Client) getInt() int {
	c.randomut.RLock()
	num := c.rando.Int()
	c.randomut.RUnlock()
	return num
}
<commit_msg>fixed timeouts<commit_after>package mqttclient

import (
	"crypto/tls"
	"errors"
	"fmt"
	mqtt "github.com/clearblade/mqtt_parsing"
	"io"
	mrand "math/rand"
	"net"
	"sync"
	"time"
)

type Client struct {

	//this is all clearblade specific
	SystemKey    string
	SystemSecret string
	AuthToken    string
	Clientid     string
	//this will usually be occupied
	//by a net.Conn
	C                   io.ReadWriteCloser
	internalOutgoingBuf chan []byte
	Timeout             time.Duration
	//thou shalt type channels consumed by others
	ClientErrorBuffer chan error

	internalErrorBuffer chan *errWrap
	shutdown_reader     chan struct{}
	shutdown_writer     chan struct{}
	//introduce a sync write mode?
	last_timeout_reccd time.Time

	shutting_down bool

	msg_store                *storage
	subscriptions            *outgoing_topics
	waiting_for_subscription *subscription_store
	//TODO:redesign around a thread-local
	//rng
	rando    *mrand.Rand
	randomut *sync.RWMutex
}

var (
	Verbose bool
)

//Start connects to the mqtt broker. It does not send the connect packet. Use the SendConnect function for that.
func (c *Client) Start(addr string, ssl *tls.Config) error {
	var con net.Conn
	var err error
	if ssl != nil {
		con, err = tls.Dial("tcp", addr, ssl)
	} else {
		con, err = net.Dial("tcp", addr)
	}
	if err != nil {
		return err
	}
	c.C = con

	go c.connectionWriter()
	go c.connectionListener()
	go c.errorTree()
	return nil
}

//NewClient allocates a new client. It is supplied with the (in order of appearance)
//Token, SystemKey,SystemSecret,Clientid, and the mqtt timeout
//Note that the following combinations of (Token|SystemKey|SystemSecret) are allowed
//(Token && SystemKey), (SystemKey && SystemSecret)
func NewClient(tok, sk, ss, cid string, timeout int) *Client {
	client := &Client{
		msg_store:                newStorage(),
		waiting_for_subscription: newSubscriptionStore(),
		subscriptions:            newOutgoingTopics(),
		SystemSecret:             ss,
		SystemKey:                sk,
		AuthToken:                tok,
		Clientid:                 cid,
		Timeout:                  time.Duration(timeout) * time.Second,
		internalOutgoingBuf:      make(chan []byte, 30),
		ClientErrorBuffer:        make(chan error, 10),
		internalErrorBuffer:      make(chan *errWrap, 2),
		shutdown_reader:          make(chan struct{}, 1),
		shutdown_writer:          make(chan struct{}, 1),
		rando:                    mrand.New(mrand.NewSource(time.Now().UnixNano())),
		randomut:                 new(sync.RWMutex),
	}
	return client
}

//sendMessage is an internal function that acts as a central point
//of failure for all of the message sending channels
//sort of like a fan-in, except this simply allows us to do
// all of the error handling logic in one place
func (c *Client) sendMessage(m mqtt.Message) error {
	select {
	case c.internalOutgoingBuf <- m.Encode():
		return nil
	default:
		return errors.New(fmt.Sprintf("Sending Channel blocked at %d", len(c.internalOutgoingBuf)))
	}
}

//connectionWriter is an internal function. it essentially sits in a goroutine and writes to the connection
//whenever it recieves a message over the channel
func (c *Client) connectionWriter() {
	if c.C == nil {
		return
	}
	for {
		select {
		case out := <-c.internalOutgoingBuf:
			_, err := c.C.Write(out)
			if err != nil {
				c.internalErrorBuffer <- &errWrap{
					err:      err,
					reciever: _CON_WRITER,
				}
				return
			}
		case <-c.shutdown_writer:
			return
		}
	}
}

//connectionListener is another function that sits on a goroutine.
//DecodePacket blocks until it reads a complete mqtt packet
func (c *Client) connectionListener() {
	if c.C == nil {
		return
	}
	mch, ech, shutdown := make(chan mqtt.Message, 10), make(chan error, 1), false
	//we have to establish an internal chain of goroutines here
	//otherwise we couldn't shutdown the listener on demand
	//since it's really hard to coordinate all the shutting down
	//when a connection drops
	//we're waiting for the connection listener to simply fail
	//this allows us to handle it a bit more gracefully
	//in order to shut down via channels directly we'd have to
	//wait for the read to fail anyway.
	go func(m chan mqtt.Message, e chan error) {
		for {
			msg, err := mqtt.DecodePacket(c.C)
			if err != nil {
				e <- err
				return
			}
			m <- msg
		}
	}(mch, ech)
	sentPing := false
	for {
		select {

		case msg := <-mch:
			//dispatch the internal functions
			c.dispatch(msg)
			sentPing = false
		case e := <-ech:
			//an error was recieved from the listenr
			if !shutdown {
				c.internalErrorBuffer <- &errWrap{
					err:      e,
					reciever: _CON_READER,
				}
			}
		case <-time.After(c.Timeout):
			//keepalive
			if !sentPing {
				c.sendMessage(&mqtt.Pingreq{})
				sentPing = true
			} else if time.Since(c.last_timeout_reccd) > c.Timeout && sentPing {
				//there was a timeout
				c.Shutdown(true)
			}

		case <-c.shutdown_reader:
			shutdown = true
			return
		}
	}
}

//this is the "business logic" of the client. it decides how each packet mutates the
//internal state of the client
func (c *Client) dispatch(msg mqtt.Message) {
	//for example we make a decision if the client sees this request or not
	//or if we have to send another message in a flow
	c.last_timeout_reccd = time.Now()
	switch msg.Type() {
	case mqtt.CONNECT:
		//shouldn't happen?
	case mqtt.CONNACK:
	case mqtt.PUBLISH:
		c.subscriptions.relay_message(msg.(*mqtt.Publish), msg.(*mqtt.Publish).Topic.Whole)
		switch msg.(*mqtt.Publish).Header.QOS {
		case 1:
			c.sendMessage(&mqtt.Puback{
				MessageId: msg.(*mqtt.Publish).MessageId,
			})
		case 2:
			c.sendMessage(&mqtt.Pubrec{
				MessageId: msg.(*mqtt.Publish).MessageId,
			})
		}
	case mqtt.PUBACK:
		//TODO:handle resend
	case mqtt.PUBREC:
		//discard, store the fact that it was recieved
		c.sendMessage(&mqtt.Pubrel{
			MessageId: msg.(*mqtt.Pubrec).MessageId,
			Header: &mqtt.StaticHeader{
				DUP:    false,
				Retain: false,
				QOS:    1,
			}})
	case mqtt.PUBREL:
		//this shouldn't have happened
	case mqtt.SUBSCRIBE:
		//this is not supposed to happen
	case mqtt.SUBACK:
		//the subscribe call blocks, so we need to forward the message
		//along that the subscribe was acknowleged so we can return
		//control flow to the parent program
		//of course, the problem is that a suback does not have
		//the subscriptions in it by name
		//but it does have the same message id
		//so we have to retrieve that and then match up the subscribe
		//TODO:NOTE THAT WE ARE ONLY USING ONE TOPIC PER SUBSCRIBE MESSSAGE
		//THIS LOGIC WILL NEED TWEAKING IF THAT CHANGES
		msg := c.msg_store.getEntry(msg.(*mqtt.Suback).MessageId)
		if msg == nil {
			//this is a bad thing to happen
			return
		}
		sub, ok := msg.(*mqtt.Subscribe)
		if !ok {
			//this is a worse thing to happen
			return
		}
		//TODO: BUG:: IF WE MAKE MULTISUBSCRIPTION, THIS WILL BREAK
		//we've now released control flow in that channel
		//also it'll allocate the userside channel and all that good stuff
		c.waiting_for_subscription.relay_message(msg, sub.Subscriptions[0].Topic.Whole)

	case mqtt.UNSUBSCRIBE:
		//this shouldn't happen
	case mqtt.UNSUBACK:
		//not gorgeous, but it do the thing
		msg := c.msg_store.getEntry(msg.(*mqtt.Unsuback).MessageId)
		unsub, ok := msg.(*mqtt.Unsubscribe)
		if !ok {
			return
		}
		c.subscriptions.remove_subscription(unsub.Topics[0].Topic.Whole)
		//TODO: do something besides prepend UNSUBSCRIBE# (the hash makes it an invalid mqtt topic, which should prevent collisions) to the front of the topic
		//but also doesn't allocate another big-ole map wrapper
		err := c.waiting_for_subscription.relay_message(msg, "UNSUBSCRIBE#"+unsub.Topics[0].Topic.Whole)
		if err != nil {
			c.ClientErrorBuffer <- err
		}
	case mqtt.PINGREQ:
		//shouldn't happen
	case mqtt.PINGRESP:
		//pingresp will reset the counter elsewhere
	case mqtt.DISCONNECT:
		//shouldn't happen
	default:
		c.ClientErrorBuffer <- fmt.Errorf("Invalid mqtt type recieved %+v", msg)
	}
}

//errorTree is the goroutine that sits on it's own goroutine and waits for
//a message to be recieved on c.internalErrorBuffer. It's our "in case of emergency break glass"
//way of reporting an error, and shutting the entire thing down
func (c *Client) errorTree() {
	//we still need to shutdown the listeners anyway
	//at least write will probably not error out

	//since you're reading the text of this fn, prepare your face for some exposition on how this mechanism works
	//so, we've spread reading and writing to the conn (or whatever) across goroutines, this is great
	//high speed low drag
	//but what happens if one goroutine encounters an error? the goroutines don't know about each other, so what do we do?
	//well, writing, and reading from to a closed connection is an error condition. so if one goroutine dies, then the other
	//will be taken down with it
	//we also use this mechanism for a regular shutdown of the client's connection, simply crashing them both and releasing the resources
	e := <-c.internalErrorBuffer
	if e.reciever != _CON_READER {
		c.shutdown_reader <- struct{}{}
	}
	if e.reciever != _CON_WRITER {
		c.shutdown_writer <- struct{}{}
	}
	if c.C != nil {
		c.C.Close()
	}
}

//Shutdown sends a disconnect packet (if asked), and then disconnects from the broker after a set time limit
func (c *Client) Shutdown(sendDisconnect bool) error {
	var err error
	if sendDisconnect {
		e := SendDisconnect(c)
		if e != nil {
			//don't return here, wait to finish the flow
			err = errors.New("While sending disconnect: " + e.Error() + "\nNote:connection was shut down anyway")
		}
		<-time.After(time.Second)
	}
	c.internalErrorBuffer <- &errWrap{reciever: _REGULAR_SHUTDOWN}
	return err
}

func (c *Client) randoPerm(i int) []int {
	c.randomut.RLock()
	order := c.rando.Perm(i)
	c.randomut.RUnlock()
	return order
}

func (c *Client) getInt() int {
	c.randomut.RLock()
	num := c.rando.Int()
	c.randomut.RUnlock()
	return num
}
<|endoftext|>


<commit_before>package quic

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/logging"
)

type client struct {
	sconn sendConn
	// If the client is created with DialAddr, we create a packet conn.
	// If it is started with Dial, we take a packet conn as a parameter.
	createdPacketConn bool

	use0RTT bool

	packetHandlers packetHandlerManager

	tlsConf *tls.Config
	config  *Config

	srcConnID  protocol.ConnectionID
	destConnID protocol.ConnectionID

	initialPacketNumber  protocol.PacketNumber
	hasNegotiatedVersion bool
	version              protocol.VersionNumber

	handshakeChan chan struct{}

	conn quicConn

	tracer    logging.ConnectionTracer
	tracingID uint64
	logger    utils.Logger
}

var (
	// make it possible to mock connection ID generation in the tests
	generateConnectionID           = protocol.GenerateConnectionID
	generateConnectionIDForInitial = protocol.GenerateConnectionIDForInitial
)

// DialAddr establishes a new QUIC connection to a server.
// It uses a new UDP connection and closes this connection when the QUIC connection is closed.
// The hostname for SNI is taken from the given address.
// The tls.Config.CipherSuites allows setting of TLS 1.3 cipher suites.
func DialAddr(
	addr string,
	tlsConf *tls.Config,
	config *Config,
) (Connection, error) {
	return DialAddrContext(context.Background(), addr, tlsConf, config)
}

// DialAddrEarly establishes a new 0-RTT QUIC connection to a server.
// It uses a new UDP connection and closes this connection when the QUIC connection is closed.
// The hostname for SNI is taken from the given address.
// The tls.Config.CipherSuites allows setting of TLS 1.3 cipher suites.
func DialAddrEarly(
	addr string,
	tlsConf *tls.Config,
	config *Config,
) (EarlyConnection, error) {
	return DialAddrEarlyContext(context.Background(), addr, tlsConf, config)
}

// DialAddrEarlyContext establishes a new 0-RTT QUIC connection to a server using provided context.
// See DialAddrEarly for details
func DialAddrEarlyContext(
	ctx context.Context,
	addr string,
	tlsConf *tls.Config,
	config *Config,
) (EarlyConnection, error) {
	conn, err := dialAddrContext(ctx, addr, tlsConf, config, true)
	if err != nil {
		return nil, err
	}
	utils.Logger.WithPrefix(utils.DefaultLogger, "client").Debugf("Returning early connection")
	return conn, nil
}

// DialAddrContext establishes a new QUIC connection to a server using the provided context.
// See DialAddr for details.
func DialAddrContext(
	ctx context.Context,
	addr string,
	tlsConf *tls.Config,
	config *Config,
) (Connection, error) {
	return dialAddrContext(ctx, addr, tlsConf, config, false)
}

func dialAddrContext(
	ctx context.Context,
	addr string,
	tlsConf *tls.Config,
	config *Config,
	use0RTT bool,
) (quicConn, error) {
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		return nil, err
	}
	return dialContext(ctx, udpConn, udpAddr, addr, tlsConf, config, use0RTT, true)
}

// Dial establishes a new QUIC connection to a server using a net.PacketConn. If
// the PacketConn satisfies the OOBCapablePacketConn interface (as a net.UDPConn
// does), ECN and packet info support will be enabled. In this case, ReadMsgUDP
// and WriteMsgUDP will be used instead of ReadFrom and WriteTo to read/write
// packets. The same PacketConn can be used for multiple calls to Dial and
// Listen, QUIC connection IDs are used for demultiplexing the different
// connections. The host parameter is used for SNI. The tls.Config must define
// an application protocol (using NextProtos).
func Dial(
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
) (Connection, error) {
	return dialContext(context.Background(), pconn, remoteAddr, host, tlsConf, config, false, false)
}

// DialEarly establishes a new 0-RTT QUIC connection to a server using a net.PacketConn.
// The same PacketConn can be used for multiple calls to Dial and Listen,
// QUIC connection IDs are used for demultiplexing the different connections.
// The host parameter is used for SNI.
// The tls.Config must define an application protocol (using NextProtos).
func DialEarly(
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
) (EarlyConnection, error) {
	return DialEarlyContext(context.Background(), pconn, remoteAddr, host, tlsConf, config)
}

// DialEarlyContext establishes a new 0-RTT QUIC connection to a server using a net.PacketConn using the provided context.
// See DialEarly for details.
func DialEarlyContext(
	ctx context.Context,
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
) (EarlyConnection, error) {
	return dialContext(ctx, pconn, remoteAddr, host, tlsConf, config, true, false)
}

// DialContext establishes a new QUIC connection to a server using a net.PacketConn using the provided context.
// See Dial for details.
func DialContext(
	ctx context.Context,
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
) (Connection, error) {
	return dialContext(ctx, pconn, remoteAddr, host, tlsConf, config, false, false)
}

func dialContext(
	ctx context.Context,
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
	use0RTT bool,
	createdPacketConn bool,
) (quicConn, error) {
	if tlsConf == nil {
		return nil, errors.New("quic: tls.Config not set")
	}
	if err := validateConfig(config); err != nil {
		return nil, err
	}
	config = populateClientConfig(config, createdPacketConn)
	packetHandlers, err := getMultiplexer().AddConn(pconn, config.ConnectionIDLength, config.StatelessResetKey, config.Tracer)
	if err != nil {
		return nil, err
	}
	c, err := newClient(pconn, remoteAddr, config, tlsConf, host, use0RTT, createdPacketConn)
	if err != nil {
		return nil, err
	}
	c.packetHandlers = packetHandlers

	c.tracingID = nextConnTracingID()
	if c.config.Tracer != nil {
		c.tracer = c.config.Tracer.TracerForConnection(
			context.WithValue(ctx, ConnectionTracingKey, c.tracingID),
			protocol.PerspectiveClient,
			c.destConnID,
		)
	}
	if c.tracer != nil {
		c.tracer.StartedConnection(c.sconn.LocalAddr(), c.sconn.RemoteAddr(), c.srcConnID, c.destConnID)
	}
	if err := c.dial(ctx); err != nil {
		return nil, err
	}
	return c.conn, nil
}

func newClient(
	pconn net.PacketConn,
	remoteAddr net.Addr,
	config *Config,
	tlsConf *tls.Config,
	host string,
	use0RTT bool,
	createdPacketConn bool,
) (*client, error) {
	if tlsConf == nil {
		tlsConf = &tls.Config{}
	}
	if tlsConf.ServerName == "" {
		sni := host
		if strings.IndexByte(sni, ':') != -1 {
			var err error
			sni, _, err = net.SplitHostPort(sni)
			if err != nil {
				return nil, err
			}
		}

		tlsConf.ServerName = sni
	}

	// check that all versions are actually supported
	if config != nil {
		for _, v := range config.Versions {
			if !protocol.IsValidVersion(v) {
				return nil, fmt.Errorf("%s is not a valid QUIC version", v)
			}
		}
	}

	srcConnID, err := generateConnectionID(config.ConnectionIDLength)
	if err != nil {
		return nil, err
	}
	destConnID, err := generateConnectionIDForInitial()
	if err != nil {
		return nil, err
	}
	c := &client{
		srcConnID:         srcConnID,
		destConnID:        destConnID,
		sconn:             newSendPconn(pconn, remoteAddr),
		createdPacketConn: createdPacketConn,
		use0RTT:           use0RTT,
		tlsConf:           tlsConf,
		config:            config,
		version:           config.Versions[0],
		handshakeChan:     make(chan struct{}),
		logger:            utils.DefaultLogger.WithPrefix("client"),
	}
	return c, nil
}

func (c *client) dial(ctx context.Context) error {
	c.logger.Infof("Starting new connection to %s (%s -> %s), source connection ID %s, destination connection ID %s, version %s", c.tlsConf.ServerName, c.sconn.LocalAddr(), c.sconn.RemoteAddr(), c.srcConnID, c.destConnID, c.version)

	c.conn = newClientConnection(
		c.sconn,
		c.packetHandlers,
		c.destConnID,
		c.srcConnID,
		c.config,
		c.tlsConf,
		c.initialPacketNumber,
		c.use0RTT,
		c.hasNegotiatedVersion,
		c.tracer,
		c.tracingID,
		c.logger,
		c.version,
	)
	c.packetHandlers.Add(c.srcConnID, c.conn)

	errorChan := make(chan error, 1)
	go func() {
		err := c.conn.run() // returns as soon as the connection is closed

		if e := (&errCloseForRecreating{}); !errors.As(err, &e) && c.createdPacketConn {
			c.packetHandlers.Destroy()
		}
		errorChan <- err
	}()

	// only set when we're using 0-RTT
	// Otherwise, earlyConnChan will be nil. Receiving from a nil chan blocks forever.
	var earlyConnChan <-chan struct{}
	if c.use0RTT {
		earlyConnChan = c.conn.earlyConnReady()
	}

	select {
	case <-ctx.Done():
		c.conn.shutdown()
		return ctx.Err()
	case err := <-errorChan:
		var recreateErr *errCloseForRecreating
		if errors.As(err, &recreateErr) {
			c.initialPacketNumber = recreateErr.nextPacketNumber
			c.version = recreateErr.nextVersion
			c.hasNegotiatedVersion = true
			return c.dial(ctx)
		}
		return err
	case <-earlyConnChan:
		// ready to send 0-RTT data
		return nil
	case <-c.conn.HandshakeComplete().Done():
		// handshake successfully completed
		return nil
	}
}
<commit_msg>clone TLS conf in newClient (#3400)<commit_after>package quic

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"

	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"github.com/lucas-clemente/quic-go/logging"
)

type client struct {
	sconn sendConn
	// If the client is created with DialAddr, we create a packet conn.
	// If it is started with Dial, we take a packet conn as a parameter.
	createdPacketConn bool

	use0RTT bool

	packetHandlers packetHandlerManager

	tlsConf *tls.Config
	config  *Config

	srcConnID  protocol.ConnectionID
	destConnID protocol.ConnectionID

	initialPacketNumber  protocol.PacketNumber
	hasNegotiatedVersion bool
	version              protocol.VersionNumber

	handshakeChan chan struct{}

	conn quicConn

	tracer    logging.ConnectionTracer
	tracingID uint64
	logger    utils.Logger
}

var (
	// make it possible to mock connection ID generation in the tests
	generateConnectionID           = protocol.GenerateConnectionID
	generateConnectionIDForInitial = protocol.GenerateConnectionIDForInitial
)

// DialAddr establishes a new QUIC connection to a server.
// It uses a new UDP connection and closes this connection when the QUIC connection is closed.
// The hostname for SNI is taken from the given address.
// The tls.Config.CipherSuites allows setting of TLS 1.3 cipher suites.
func DialAddr(
	addr string,
	tlsConf *tls.Config,
	config *Config,
) (Connection, error) {
	return DialAddrContext(context.Background(), addr, tlsConf, config)
}

// DialAddrEarly establishes a new 0-RTT QUIC connection to a server.
// It uses a new UDP connection and closes this connection when the QUIC connection is closed.
// The hostname for SNI is taken from the given address.
// The tls.Config.CipherSuites allows setting of TLS 1.3 cipher suites.
func DialAddrEarly(
	addr string,
	tlsConf *tls.Config,
	config *Config,
) (EarlyConnection, error) {
	return DialAddrEarlyContext(context.Background(), addr, tlsConf, config)
}

// DialAddrEarlyContext establishes a new 0-RTT QUIC connection to a server using provided context.
// See DialAddrEarly for details
func DialAddrEarlyContext(
	ctx context.Context,
	addr string,
	tlsConf *tls.Config,
	config *Config,
) (EarlyConnection, error) {
	conn, err := dialAddrContext(ctx, addr, tlsConf, config, true)
	if err != nil {
		return nil, err
	}
	utils.Logger.WithPrefix(utils.DefaultLogger, "client").Debugf("Returning early connection")
	return conn, nil
}

// DialAddrContext establishes a new QUIC connection to a server using the provided context.
// See DialAddr for details.
func DialAddrContext(
	ctx context.Context,
	addr string,
	tlsConf *tls.Config,
	config *Config,
) (Connection, error) {
	return dialAddrContext(ctx, addr, tlsConf, config, false)
}

func dialAddrContext(
	ctx context.Context,
	addr string,
	tlsConf *tls.Config,
	config *Config,
	use0RTT bool,
) (quicConn, error) {
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}
	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		return nil, err
	}
	return dialContext(ctx, udpConn, udpAddr, addr, tlsConf, config, use0RTT, true)
}

// Dial establishes a new QUIC connection to a server using a net.PacketConn. If
// the PacketConn satisfies the OOBCapablePacketConn interface (as a net.UDPConn
// does), ECN and packet info support will be enabled. In this case, ReadMsgUDP
// and WriteMsgUDP will be used instead of ReadFrom and WriteTo to read/write
// packets. The same PacketConn can be used for multiple calls to Dial and
// Listen, QUIC connection IDs are used for demultiplexing the different
// connections. The host parameter is used for SNI. The tls.Config must define
// an application protocol (using NextProtos).
func Dial(
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
) (Connection, error) {
	return dialContext(context.Background(), pconn, remoteAddr, host, tlsConf, config, false, false)
}

// DialEarly establishes a new 0-RTT QUIC connection to a server using a net.PacketConn.
// The same PacketConn can be used for multiple calls to Dial and Listen,
// QUIC connection IDs are used for demultiplexing the different connections.
// The host parameter is used for SNI.
// The tls.Config must define an application protocol (using NextProtos).
func DialEarly(
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
) (EarlyConnection, error) {
	return DialEarlyContext(context.Background(), pconn, remoteAddr, host, tlsConf, config)
}

// DialEarlyContext establishes a new 0-RTT QUIC connection to a server using a net.PacketConn using the provided context.
// See DialEarly for details.
func DialEarlyContext(
	ctx context.Context,
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
) (EarlyConnection, error) {
	return dialContext(ctx, pconn, remoteAddr, host, tlsConf, config, true, false)
}

// DialContext establishes a new QUIC connection to a server using a net.PacketConn using the provided context.
// See Dial for details.
func DialContext(
	ctx context.Context,
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
) (Connection, error) {
	return dialContext(ctx, pconn, remoteAddr, host, tlsConf, config, false, false)
}

func dialContext(
	ctx context.Context,
	pconn net.PacketConn,
	remoteAddr net.Addr,
	host string,
	tlsConf *tls.Config,
	config *Config,
	use0RTT bool,
	createdPacketConn bool,
) (quicConn, error) {
	if tlsConf == nil {
		return nil, errors.New("quic: tls.Config not set")
	}
	if err := validateConfig(config); err != nil {
		return nil, err
	}
	config = populateClientConfig(config, createdPacketConn)
	packetHandlers, err := getMultiplexer().AddConn(pconn, config.ConnectionIDLength, config.StatelessResetKey, config.Tracer)
	if err != nil {
		return nil, err
	}
	c, err := newClient(pconn, remoteAddr, config, tlsConf, host, use0RTT, createdPacketConn)
	if err != nil {
		return nil, err
	}
	c.packetHandlers = packetHandlers

	c.tracingID = nextConnTracingID()
	if c.config.Tracer != nil {
		c.tracer = c.config.Tracer.TracerForConnection(
			context.WithValue(ctx, ConnectionTracingKey, c.tracingID),
			protocol.PerspectiveClient,
			c.destConnID,
		)
	}
	if c.tracer != nil {
		c.tracer.StartedConnection(c.sconn.LocalAddr(), c.sconn.RemoteAddr(), c.srcConnID, c.destConnID)
	}
	if err := c.dial(ctx); err != nil {
		return nil, err
	}
	return c.conn, nil
}

func newClient(
	pconn net.PacketConn,
	remoteAddr net.Addr,
	config *Config,
	tlsConf *tls.Config,
	host string,
	use0RTT bool,
	createdPacketConn bool,
) (*client, error) {
	if tlsConf == nil {
		tlsConf = &tls.Config{}
	} else {
		tlsConf = tlsConf.Clone()
	}
	if tlsConf.ServerName == "" {
		sni := host
		if strings.IndexByte(sni, ':') != -1 {
			var err error
			sni, _, err = net.SplitHostPort(sni)
			if err != nil {
				return nil, err
			}
		}

		tlsConf.ServerName = sni
	}

	// check that all versions are actually supported
	if config != nil {
		for _, v := range config.Versions {
			if !protocol.IsValidVersion(v) {
				return nil, fmt.Errorf("%s is not a valid QUIC version", v)
			}
		}
	}

	srcConnID, err := generateConnectionID(config.ConnectionIDLength)
	if err != nil {
		return nil, err
	}
	destConnID, err := generateConnectionIDForInitial()
	if err != nil {
		return nil, err
	}
	c := &client{
		srcConnID:         srcConnID,
		destConnID:        destConnID,
		sconn:             newSendPconn(pconn, remoteAddr),
		createdPacketConn: createdPacketConn,
		use0RTT:           use0RTT,
		tlsConf:           tlsConf,
		config:            config,
		version:           config.Versions[0],
		handshakeChan:     make(chan struct{}),
		logger:            utils.DefaultLogger.WithPrefix("client"),
	}
	return c, nil
}

func (c *client) dial(ctx context.Context) error {
	c.logger.Infof("Starting new connection to %s (%s -> %s), source connection ID %s, destination connection ID %s, version %s", c.tlsConf.ServerName, c.sconn.LocalAddr(), c.sconn.RemoteAddr(), c.srcConnID, c.destConnID, c.version)

	c.conn = newClientConnection(
		c.sconn,
		c.packetHandlers,
		c.destConnID,
		c.srcConnID,
		c.config,
		c.tlsConf,
		c.initialPacketNumber,
		c.use0RTT,
		c.hasNegotiatedVersion,
		c.tracer,
		c.tracingID,
		c.logger,
		c.version,
	)
	c.packetHandlers.Add(c.srcConnID, c.conn)

	errorChan := make(chan error, 1)
	go func() {
		err := c.conn.run() // returns as soon as the connection is closed

		if e := (&errCloseForRecreating{}); !errors.As(err, &e) && c.createdPacketConn {
			c.packetHandlers.Destroy()
		}
		errorChan <- err
	}()

	// only set when we're using 0-RTT
	// Otherwise, earlyConnChan will be nil. Receiving from a nil chan blocks forever.
	var earlyConnChan <-chan struct{}
	if c.use0RTT {
		earlyConnChan = c.conn.earlyConnReady()
	}

	select {
	case <-ctx.Done():
		c.conn.shutdown()
		return ctx.Err()
	case err := <-errorChan:
		var recreateErr *errCloseForRecreating
		if errors.As(err, &recreateErr) {
			c.initialPacketNumber = recreateErr.nextPacketNumber
			c.version = recreateErr.nextVersion
			c.hasNegotiatedVersion = true
			return c.dial(ctx)
		}
		return err
	case <-earlyConnChan:
		// ready to send 0-RTT data
		return nil
	case <-c.conn.HandshakeComplete().Done():
		// handshake successfully completed
		return nil
	}
}
<|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
	Color         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(),
		Color:  RandomColor(),
		Msg:    make(chan string, MSG_BUFFER),
		ready:  make(chan struct{}, 1),
	}
}

func (c *Client) ColoredName() string {
    return ColorString(c.Color, c.Name)
}

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] ", c.ColoredName()))
}

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 "/test-colors": // Shh, this command is a secret!
				c.Write(ColorString("32", "Lorem ipsum dolor sit amet,"))
				c.Write("consectetur " + ColorString("31;1", "adipiscing") + " elit.")
			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.ColoredName(), 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 := RE_STRIP_TEXT.ReplaceAllString(string(client.Conn.ClientVersion()), "")
						if len(version) > 100 {
							version = "Evil Jerk with a superlong string"
						}
						c.Msg <- fmt.Sprintf("-> %s is %s via %s", client.ColoredName(), 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.ColoredName()))
						c.Server.Ban(fingerprint, nil)
						client.Conn.Close()
						c.Server.Broadcast(fmt.Sprintf("* %s was banned by %s", parts[1], c.ColoredName()), 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.ColoredName()))
						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.ColoredName()))
					}
				}
			default:
				c.Msg <- fmt.Sprintf("-> Invalid command: %s", line)
			}
			continue
		}

		msg := fmt.Sprintf("%s: %s", c.ColoredName(), 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.ColoredName())

	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>Block the broadcast of an empty message<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
	Color         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(),
		Color:  RandomColor(),
		Msg:    make(chan string, MSG_BUFFER),
		ready:  make(chan struct{}, 1),
	}
}

func (c *Client) ColoredName() string {
    return ColorString(c.Color, c.Name)
}

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] ", c.ColoredName()))
}

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 "/test-colors": // Shh, this command is a secret!
				c.Write(ColorString("32", "Lorem ipsum dolor sit amet,"))
				c.Write("consectetur " + ColorString("31;1", "adipiscing") + " elit.")
			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.ColoredName(), 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 := RE_STRIP_TEXT.ReplaceAllString(string(client.Conn.ClientVersion()), "")
						if len(version) > 100 {
							version = "Evil Jerk with a superlong string"
						}
						c.Msg <- fmt.Sprintf("-> %s is %s via %s", client.ColoredName(), 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.ColoredName()))
						c.Server.Ban(fingerprint, nil)
						client.Conn.Close()
						c.Server.Broadcast(fmt.Sprintf("* %s was banned by %s", parts[1], c.ColoredName()), 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.ColoredName()))
						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.ColoredName()))
					}
				}
			default:
				c.Msg <- fmt.Sprintf("-> Invalid command: %s", line)
			}
			continue
		}

		msg := fmt.Sprintf("%s: %s", c.ColoredName(), line)
		if c.IsSilenced() || len(msg) > 1000 || len(line) < 1 {
			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.ColoredName())

	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|>