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

package gomatrixserverlib

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/matrix-org/gomatrix"
	"github.com/matrix-org/util"
	"github.com/sirupsen/logrus"
)

// Default HTTPS request timeout
const requestTimeout time.Duration = time.Duration(30) * time.Second

// A Client makes request to the federation listeners of matrix
// homeservers
type Client struct {
	client http.Client
}

// UserInfo represents information about a user.
type UserInfo struct {
	Sub string `json:"sub"`
}

// NewClient makes a new Client (with default timeout)
func NewClient() *Client {
	return NewClientWithTimeout(requestTimeout, newFederationTripper())
}

// NewClientWithTransport makes a new Client with an existing transport
func NewClientWithTransport(transport http.RoundTripper) *Client {
	return NewClientWithTimeout(requestTimeout, transport)
}

// NewClientWithTimeout makes a new Client with a specified request timeout
func NewClientWithTimeout(timeout time.Duration, transport http.RoundTripper) *Client {
	return &Client{
		client: http.Client{
			Transport: transport,
			Timeout:   timeout,
		},
	}
}

type federationTripper struct {
	// transports maps an TLS server name with an HTTP transport.
	transports      map[string]http.RoundTripper
	transportsMutex sync.Mutex
}

func newFederationTripper() *federationTripper {
	return &federationTripper{
		transports: make(map[string]http.RoundTripper),
	}
}

// getTransport returns a http.Transport instance with a TLS configuration using
// the given server name for SNI. It also creates the instance if there isn't
// any for this server name.
// We need to use one transport per TLS server name (instead of giving our round
// tripper a single transport) because there is no way to specify the TLS
// ServerName on a per-connection basis.
func (f *federationTripper) getTransport(tlsServerName string) (transport http.RoundTripper) {
	var ok bool

	f.transportsMutex.Lock()

	// Create the transport if we don't have any for this TLS server name.
	if transport, ok = f.transports[tlsServerName]; !ok {
		transport = &http.Transport{
			TLSClientConfig: &tls.Config{
				ServerName: tlsServerName,
				// TODO: Remove this when we enforce MSC1711.
				InsecureSkipVerify: true,
			},
		}

		f.transports[tlsServerName] = transport
	}

	f.transportsMutex.Unlock()

	return transport
}

func makeHTTPSURL(u *url.URL, addr string) (httpsURL url.URL) {
	httpsURL = *u
	httpsURL.Scheme = "https"
	httpsURL.Host = addr
	return
}

func (f *federationTripper) RoundTrip(r *http.Request) (*http.Response, error) {
	serverName := ServerName(r.URL.Host)
	resolutionResults, err := ResolveServer(serverName)
	if err != nil {
		return nil, err
	}

	if len(resolutionResults) == 0 {
		return nil, fmt.Errorf("no address found for matrix host %v", serverName)
	}

	var resp *http.Response
	// TODO: respect the priority and weight fields from the SRV record
	for _, result := range resolutionResults {
		u := makeHTTPSURL(r.URL, result.Destination)
		r.URL = &u
		r.Host = string(result.Host)
		resp, err = f.getTransport(result.TLSServerName).RoundTrip(r)
		if err == nil {
			return resp, nil
		}
		util.GetLogger(r.Context()).Warnf("Error sending request to %s: %v",
			u.String(), err)
	}

	// just return the most recent error
	return nil, err
}

// LookupUserInfo gets information about a user from a given matrix homeserver
// using a bearer access token.
func (fc *Client) LookupUserInfo(
	ctx context.Context, matrixServer ServerName, token string,
) (u UserInfo, err error) {
	url := url.URL{
		Scheme:   "matrix",
		Host:     string(matrixServer),
		Path:     "/_matrix/federation/v1/openid/userinfo",
		RawQuery: url.Values{"access_token": []string{token}}.Encode(),
	}

	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		return
	}

	var response *http.Response
	response, err = fc.DoHTTPRequest(ctx, req)
	if response != nil {
		defer response.Body.Close() // nolint: errcheck
	}
	if err != nil {
		return
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		var errorOutput []byte
		errorOutput, err = ioutil.ReadAll(response.Body)
		if err != nil {
			return
		}
		err = fmt.Errorf("HTTP %d : %s", response.StatusCode, errorOutput)
		return
	}

	err = json.NewDecoder(response.Body).Decode(&u)
	if err != nil {
		return
	}

	userParts := strings.SplitN(u.Sub, ":", 2)
	if len(userParts) != 2 || userParts[1] != string(matrixServer) {
		err = fmt.Errorf("userID doesn't match server name '%v' != '%v'", u.Sub, matrixServer)
		return
	}

	return
}

// GetServerKeys asks a matrix server for its signing keys and TLS cert
func (fc *Client) GetServerKeys(
	ctx context.Context, matrixServer ServerName,
) (ServerKeys, error) {
	url := url.URL{
		Scheme: "matrix",
		Host:   string(matrixServer),
		Path:   "/_matrix/key/v2/server",
	}

	var body ServerKeys
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		return body, err
	}

	err = fc.DoRequestAndParseResponse(
		ctx, req, &body,
	)
	return body, err
}

// GetVersion gets the version information of a homeserver.
// See https://matrix.org/docs/spec/server_server/r0.1.1.html#get-matrix-federation-v1-version
func (fc *Client) GetVersion(
	ctx context.Context, s ServerName,
) (res Version, err error) {
	// Construct a request for version information
	url := url.URL{
		Scheme: "matrix",
		Host:   string(s),
		Path:   "/_matrix/federation/v1/version",
	}
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		return
	}

	// Make the request and parse the response
	err = fc.DoRequestAndParseResponse(ctx, req, &res)
	return
}

// LookupServerKeys looks up the keys for a matrix server from a matrix server.
// The first argument is the name of the matrix server to download the keys from.
// The second argument is a map from (server name, key ID) pairs to timestamps.
// The (server name, key ID) pair identifies the key to download.
// The timestamps tell the server when the keys need to be valid until.
// Perspective servers can use that timestamp to determine whether they can
// return a cached copy of the keys or whether they will need to retrieve a fresh
// copy of the keys.
// Returns the keys returned by the server, or an error if there was a problem talking to the server.
func (fc *Client) LookupServerKeys(
	ctx context.Context, matrixServer ServerName, keyRequests map[PublicKeyLookupRequest]Timestamp,
) ([]ServerKeys, error) {
	url := url.URL{
		Scheme: "matrix",
		Host:   string(matrixServer),
		Path:   "/_matrix/key/v2/query",
	}

	// The request format is:
	// { "server_keys": { "<server_name>": { "<key_id>": { "minimum_valid_until_ts": <ts> }}}
	type keyreq struct {
		MinimumValidUntilTS Timestamp `json:"minimum_valid_until_ts"`
	}
	request := struct {
		ServerKeyMap map[ServerName]map[KeyID]keyreq `json:"server_keys"`
	}{map[ServerName]map[KeyID]keyreq{}}
	for k, ts := range keyRequests {
		server := request.ServerKeyMap[k.ServerName]
		if server == nil {
			server = map[KeyID]keyreq{}
			request.ServerKeyMap[k.ServerName] = server
		}
		server[k.KeyID] = keyreq{ts}
	}

	requestBytes, err := json.Marshal(request)
	if err != nil {
		return nil, err
	}

	var body struct {
		ServerKeyList []ServerKeys `json:"server_keys"`
	}

	req, err := http.NewRequest("POST", url.String(), bytes.NewBuffer(requestBytes))
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/json")

	err = fc.DoRequestAndParseResponse(
		ctx, req, &body,
	)
	if err != nil {
		return nil, err
	}

	return body.ServerKeyList, nil
}

// CreateMediaDownloadRequest creates a request for media on a homeserver and returns the http.Response or an error
func (fc *Client) CreateMediaDownloadRequest(
	ctx context.Context, matrixServer ServerName, mediaID string,
) (*http.Response, error) {
	requestURL := "matrix://" + string(matrixServer) + "/_matrix/media/v1/download/" + string(matrixServer) + "/" + mediaID
	req, err := http.NewRequest("GET", requestURL, nil)
	if err != nil {
		return nil, err
	}

	return fc.DoHTTPRequest(ctx, req)
}

// DoRequestAndParseResponse calls DoHTTPRequest and then decodes the response.
//
// If the HTTP response is not a 200, an attempt is made to parse the response
// body into a gomatrix.RespError. In any case, a non-200 response will result
// in a gomatrix.HTTPError.
//
func (fc *Client) DoRequestAndParseResponse(
	ctx context.Context,
	req *http.Request,
	result interface{},
) error {
	response, err := fc.DoHTTPRequest(ctx, req)
	if response != nil {
		defer response.Body.Close() // nolint: errcheck
	}
	if err != nil {
		return err
	}

	if response.StatusCode/100 != 2 { // not 2xx
		// Adapted from https://github.com/matrix-org/gomatrix/blob/master/client.go
		var contents []byte
		contents, err = ioutil.ReadAll(response.Body)
		if err != nil {
			return err
		}

		var wrap error
		var respErr gomatrix.RespError
		if _ = json.Unmarshal(contents, &respErr); respErr.ErrCode != "" {
			wrap = respErr
		}

		// If we failed to decode as RespError, don't just drop the HTTP body, include it in the
		// HTTP error instead (e.g proxy errors which return HTML).
		msg := fmt.Sprintf("Failed to %s JSON (hostname %q path %q)", req.Method, req.Host, req.URL.Path)
		if wrap == nil {
			msg += ": " + string(contents)
		}

		return gomatrix.HTTPError{
			Code:         response.StatusCode,
			Message:      msg,
			WrappedError: wrap,
		}
	}

	if err = json.NewDecoder(response.Body).Decode(result); err != nil {
		return err
	}

	return nil
}

// DoHTTPRequest creates an outgoing request ID and adds it to the context
// before sending off the request and awaiting a response.
//
// If the returned error is nil, the Response will contain a non-nil
// Body which the caller is expected to close.
//
func (fc *Client) DoHTTPRequest(ctx context.Context, req *http.Request) (*http.Response, error) {
	reqID := util.RandomString(12)
	logger := util.GetLogger(ctx).WithFields(logrus.Fields{
		"out.req.ID":     reqID,
		"out.req.method": req.Method,
		"out.req.uri":    req.URL,
	})
	logger.Info("Outgoing request")
	newCtx := util.ContextWithLogger(ctx, logger)

	start := time.Now()
	resp, err := fc.client.Do(req.WithContext(newCtx))
	if err != nil {
		logger.WithField("error", err).Warn("Outgoing request failed")
		return nil, err
	}

	// we haven't yet read the body, so this is slightly premature, but it's the easiest place.
	logger.WithFields(logrus.Fields{
		"out.req.code":        resp.StatusCode,
		"out.req.duration_ms": int(time.Since(start) / time.Millisecond),
	}).Info("Outgoing request returned")

	return resp, nil
}
<commit_msg>Set allow_remote=false for loop avoidance on media fetches<commit_after>/* Copyright 2016-2017 Vector Creations Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package gomatrixserverlib

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/matrix-org/gomatrix"
	"github.com/matrix-org/util"
	"github.com/sirupsen/logrus"
)

// Default HTTPS request timeout
const requestTimeout time.Duration = time.Duration(30) * time.Second

// A Client makes request to the federation listeners of matrix
// homeservers
type Client struct {
	client http.Client
}

// UserInfo represents information about a user.
type UserInfo struct {
	Sub string `json:"sub"`
}

// NewClient makes a new Client (with default timeout)
func NewClient() *Client {
	return NewClientWithTimeout(requestTimeout, newFederationTripper())
}

// NewClientWithTransport makes a new Client with an existing transport
func NewClientWithTransport(transport http.RoundTripper) *Client {
	return NewClientWithTimeout(requestTimeout, transport)
}

// NewClientWithTimeout makes a new Client with a specified request timeout
func NewClientWithTimeout(timeout time.Duration, transport http.RoundTripper) *Client {
	return &Client{
		client: http.Client{
			Transport: transport,
			Timeout:   timeout,
		},
	}
}

type federationTripper struct {
	// transports maps an TLS server name with an HTTP transport.
	transports      map[string]http.RoundTripper
	transportsMutex sync.Mutex
}

func newFederationTripper() *federationTripper {
	return &federationTripper{
		transports: make(map[string]http.RoundTripper),
	}
}

// getTransport returns a http.Transport instance with a TLS configuration using
// the given server name for SNI. It also creates the instance if there isn't
// any for this server name.
// We need to use one transport per TLS server name (instead of giving our round
// tripper a single transport) because there is no way to specify the TLS
// ServerName on a per-connection basis.
func (f *federationTripper) getTransport(tlsServerName string) (transport http.RoundTripper) {
	var ok bool

	f.transportsMutex.Lock()

	// Create the transport if we don't have any for this TLS server name.
	if transport, ok = f.transports[tlsServerName]; !ok {
		transport = &http.Transport{
			TLSClientConfig: &tls.Config{
				ServerName: tlsServerName,
				// TODO: Remove this when we enforce MSC1711.
				InsecureSkipVerify: true,
			},
		}

		f.transports[tlsServerName] = transport
	}

	f.transportsMutex.Unlock()

	return transport
}

func makeHTTPSURL(u *url.URL, addr string) (httpsURL url.URL) {
	httpsURL = *u
	httpsURL.Scheme = "https"
	httpsURL.Host = addr
	return
}

func (f *federationTripper) RoundTrip(r *http.Request) (*http.Response, error) {
	serverName := ServerName(r.URL.Host)
	resolutionResults, err := ResolveServer(serverName)
	if err != nil {
		return nil, err
	}

	if len(resolutionResults) == 0 {
		return nil, fmt.Errorf("no address found for matrix host %v", serverName)
	}

	var resp *http.Response
	// TODO: respect the priority and weight fields from the SRV record
	for _, result := range resolutionResults {
		u := makeHTTPSURL(r.URL, result.Destination)
		r.URL = &u
		r.Host = string(result.Host)
		resp, err = f.getTransport(result.TLSServerName).RoundTrip(r)
		if err == nil {
			return resp, nil
		}
		util.GetLogger(r.Context()).Warnf("Error sending request to %s: %v",
			u.String(), err)
	}

	// just return the most recent error
	return nil, err
}

// LookupUserInfo gets information about a user from a given matrix homeserver
// using a bearer access token.
func (fc *Client) LookupUserInfo(
	ctx context.Context, matrixServer ServerName, token string,
) (u UserInfo, err error) {
	url := url.URL{
		Scheme:   "matrix",
		Host:     string(matrixServer),
		Path:     "/_matrix/federation/v1/openid/userinfo",
		RawQuery: url.Values{"access_token": []string{token}}.Encode(),
	}

	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		return
	}

	var response *http.Response
	response, err = fc.DoHTTPRequest(ctx, req)
	if response != nil {
		defer response.Body.Close() // nolint: errcheck
	}
	if err != nil {
		return
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		var errorOutput []byte
		errorOutput, err = ioutil.ReadAll(response.Body)
		if err != nil {
			return
		}
		err = fmt.Errorf("HTTP %d : %s", response.StatusCode, errorOutput)
		return
	}

	err = json.NewDecoder(response.Body).Decode(&u)
	if err != nil {
		return
	}

	userParts := strings.SplitN(u.Sub, ":", 2)
	if len(userParts) != 2 || userParts[1] != string(matrixServer) {
		err = fmt.Errorf("userID doesn't match server name '%v' != '%v'", u.Sub, matrixServer)
		return
	}

	return
}

// GetServerKeys asks a matrix server for its signing keys and TLS cert
func (fc *Client) GetServerKeys(
	ctx context.Context, matrixServer ServerName,
) (ServerKeys, error) {
	url := url.URL{
		Scheme: "matrix",
		Host:   string(matrixServer),
		Path:   "/_matrix/key/v2/server",
	}

	var body ServerKeys
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		return body, err
	}

	err = fc.DoRequestAndParseResponse(
		ctx, req, &body,
	)
	return body, err
}

// GetVersion gets the version information of a homeserver.
// See https://matrix.org/docs/spec/server_server/r0.1.1.html#get-matrix-federation-v1-version
func (fc *Client) GetVersion(
	ctx context.Context, s ServerName,
) (res Version, err error) {
	// Construct a request for version information
	url := url.URL{
		Scheme: "matrix",
		Host:   string(s),
		Path:   "/_matrix/federation/v1/version",
	}
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		return
	}

	// Make the request and parse the response
	err = fc.DoRequestAndParseResponse(ctx, req, &res)
	return
}

// LookupServerKeys looks up the keys for a matrix server from a matrix server.
// The first argument is the name of the matrix server to download the keys from.
// The second argument is a map from (server name, key ID) pairs to timestamps.
// The (server name, key ID) pair identifies the key to download.
// The timestamps tell the server when the keys need to be valid until.
// Perspective servers can use that timestamp to determine whether they can
// return a cached copy of the keys or whether they will need to retrieve a fresh
// copy of the keys.
// Returns the keys returned by the server, or an error if there was a problem talking to the server.
func (fc *Client) LookupServerKeys(
	ctx context.Context, matrixServer ServerName, keyRequests map[PublicKeyLookupRequest]Timestamp,
) ([]ServerKeys, error) {
	url := url.URL{
		Scheme: "matrix",
		Host:   string(matrixServer),
		Path:   "/_matrix/key/v2/query",
	}

	// The request format is:
	// { "server_keys": { "<server_name>": { "<key_id>": { "minimum_valid_until_ts": <ts> }}}
	type keyreq struct {
		MinimumValidUntilTS Timestamp `json:"minimum_valid_until_ts"`
	}
	request := struct {
		ServerKeyMap map[ServerName]map[KeyID]keyreq `json:"server_keys"`
	}{map[ServerName]map[KeyID]keyreq{}}
	for k, ts := range keyRequests {
		server := request.ServerKeyMap[k.ServerName]
		if server == nil {
			server = map[KeyID]keyreq{}
			request.ServerKeyMap[k.ServerName] = server
		}
		server[k.KeyID] = keyreq{ts}
	}

	requestBytes, err := json.Marshal(request)
	if err != nil {
		return nil, err
	}

	var body struct {
		ServerKeyList []ServerKeys `json:"server_keys"`
	}

	req, err := http.NewRequest("POST", url.String(), bytes.NewBuffer(requestBytes))
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/json")

	err = fc.DoRequestAndParseResponse(
		ctx, req, &body,
	)
	if err != nil {
		return nil, err
	}

	return body.ServerKeyList, nil
}

// CreateMediaDownloadRequest creates a request for media on a homeserver and returns the http.Response or an error
func (fc *Client) CreateMediaDownloadRequest(
	ctx context.Context, matrixServer ServerName, mediaID string,
) (*http.Response, error) {
	requestURL := "matrix://" + string(matrixServer) + "/_matrix/media/v1/download/" + string(matrixServer) + "/" + mediaID + "?allow_remote=false"
	req, err := http.NewRequest("GET", requestURL, nil)
	if err != nil {
		return nil, err
	}

	return fc.DoHTTPRequest(ctx, req)
}

// DoRequestAndParseResponse calls DoHTTPRequest and then decodes the response.
//
// If the HTTP response is not a 200, an attempt is made to parse the response
// body into a gomatrix.RespError. In any case, a non-200 response will result
// in a gomatrix.HTTPError.
//
func (fc *Client) DoRequestAndParseResponse(
	ctx context.Context,
	req *http.Request,
	result interface{},
) error {
	response, err := fc.DoHTTPRequest(ctx, req)
	if response != nil {
		defer response.Body.Close() // nolint: errcheck
	}
	if err != nil {
		return err
	}

	if response.StatusCode/100 != 2 { // not 2xx
		// Adapted from https://github.com/matrix-org/gomatrix/blob/master/client.go
		var contents []byte
		contents, err = ioutil.ReadAll(response.Body)
		if err != nil {
			return err
		}

		var wrap error
		var respErr gomatrix.RespError
		if _ = json.Unmarshal(contents, &respErr); respErr.ErrCode != "" {
			wrap = respErr
		}

		// If we failed to decode as RespError, don't just drop the HTTP body, include it in the
		// HTTP error instead (e.g proxy errors which return HTML).
		msg := fmt.Sprintf("Failed to %s JSON (hostname %q path %q)", req.Method, req.Host, req.URL.Path)
		if wrap == nil {
			msg += ": " + string(contents)
		}

		return gomatrix.HTTPError{
			Code:         response.StatusCode,
			Message:      msg,
			WrappedError: wrap,
		}
	}

	if err = json.NewDecoder(response.Body).Decode(result); err != nil {
		return err
	}

	return nil
}

// DoHTTPRequest creates an outgoing request ID and adds it to the context
// before sending off the request and awaiting a response.
//
// If the returned error is nil, the Response will contain a non-nil
// Body which the caller is expected to close.
//
func (fc *Client) DoHTTPRequest(ctx context.Context, req *http.Request) (*http.Response, error) {
	reqID := util.RandomString(12)
	logger := util.GetLogger(ctx).WithFields(logrus.Fields{
		"out.req.ID":     reqID,
		"out.req.method": req.Method,
		"out.req.uri":    req.URL,
	})
	logger.Info("Outgoing request")
	newCtx := util.ContextWithLogger(ctx, logger)

	start := time.Now()
	resp, err := fc.client.Do(req.WithContext(newCtx))
	if err != nil {
		logger.WithField("error", err).Warn("Outgoing request failed")
		return nil, err
	}

	// we haven't yet read the body, so this is slightly premature, but it's the easiest place.
	logger.WithFields(logrus.Fields{
		"out.req.code":        resp.StatusCode,
		"out.req.duration_ms": int(time.Since(start) / time.Millisecond),
	}).Info("Outgoing request returned")

	return resp, nil
}
<|endoftext|>


<commit_before>package synapse

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/philhofer/fwd"
	"github.com/tinylib/msgp/msgp"
)

const (
	defaultTimeout = 3000 // three seconds

	// waiter "high water mark"
	// TODO(maybe): make this adjustable.
	waiterHWM = 32
)

const (
	clientClosed = iota
	clientOpen
)

var (
	// ErrClosed is returns when a call is attempted
	// on a closed client.
	ErrClosed = errors.New("synapse: client is closed")

	// ErrTimeout is returned when a server
	// doesn't respond to a request before
	// the client's timeout scavenger can
	// free the waiting goroutine
	ErrTimeout = errors.New("synapse: the server didn't respond in time")

	// ErrTooLarge is returned when the message
	// size is larger than 65,535 bytes.
	ErrTooLarge = errors.New("synapse: message body too large")
)

// Dial creates a new client by dialing
// the provided network and remote address.
// The provided timeout is used as the timeout
// for requests, in milliseconds.
func Dial(network string, raddr string, timeout time.Duration) (*Client, error) {
	conn, err := net.Dial(network, raddr)
	if err != nil {
		return nil, err
	}
	return NewClient(conn, timeout)
}

// DialTLS acts identically to Dial, except that it dials the connection
// over TLS using the provided *tls.Config.
func DialTLS(network, raddr string, timeout time.Duration, config *tls.Config) (*Client, error) {
	conn, err := tls.Dial(network, raddr, config)
	if err != nil {
		return nil, err
	}
	return NewClient(conn, timeout)
}

// NewClient creates a new client from an
// existing net.Conn. Timeout is the maximum time,
// in milliseconds, to wait for server responses
// before sending an error to the caller. NewClient
// fails with an error if it cannot ping the server
// over the connection.
func NewClient(c net.Conn, timeout time.Duration) (*Client, error) {
	cl := &Client{
		conn:    c,
		writing: make(chan *waiter, waiterHWM),
		done:    make(chan struct{}),
		state:   clientOpen,
	}
	go cl.readLoop()
	go cl.writeLoop()

	if timeout <= 0 {
		timeout = defaultTimeout
	}

	go cl.timeoutLoop(timeout)

	// do a ping to check
	// for sanity
	err := cl.ping()
	if err != nil {
		cl.Close()
		return nil, fmt.Errorf("synapse: ping failed: %s", err)
	}

	return cl, nil
}

// Client is a client to
// a single synapse server.
type Client struct {
	conn    net.Conn       // connection
	csn     uint64         // sequence number; atomic
	writing chan *waiter   // queue to write to conn; size is effectively HWM
	done    chan struct{}  // closed during (*Client).Close to shut down timeoutloop
	wg      sync.WaitGroup // outstanding client procs
	state   uint32         // open, closed, etc.
	pending wMap           // map seq number to waiting handler
}

// used to transfer control
// flow to blocking goroutines
type waiter struct {
	next   *waiter    // next in linked list, or nil
	parent *Client    // parent *client
	seq    uint64     // sequence number
	done   sync.Mutex // for notifying response; locked is default
	err    error      // response error on wakeup, if applicable
	in     []byte     // response body
	reap   bool       // can reap for timeout
	static bool       // is part of the statically allocated arena
}

// Close idempotently closes the
// client's connection to the server.
// Close returns once every waiting
// goroutine has either received a
// response or timed out; goroutines
// with requests in progress are not interrupted.
func (c *Client) Close() error {
	// there can only be one winner of the Close() race
	if !atomic.CompareAndSwapUint32(&c.state, clientOpen, clientClosed) {
		return ErrClosed
	}

	c.wg.Wait()
	close(c.done)
	close(c.writing)
	return c.conn.Close()
}

// close with error
// sets the status of every waiting
// goroutine to 'err' and unblocks it.
func (c *Client) closeError(err error) {
	if !atomic.CompareAndSwapUint32(&c.state, clientOpen, clientClosed) {
		return
	}

	err = fmt.Errorf("synapse: fatal error: %s", err)

	// we can't actually guarantee that we will preempt
	// every goroutine, but we can try.
	c.pending.flush(err)
	for c.pending.length() > 0 {
		c.pending.flush(err)
	}

	c.wg.Wait()
	close(c.done)
	close(c.writing)
	c.conn.Close()
}

// a handler for io.Read() and io.Write(),
// e.g.
//  if !c.do(w.Write(data)) { goto fail }
func (c *Client) do(_ int, err error) bool {
	if err != nil {
		c.closeError(err)
		return false
	}
	return true
}

// readLoop continuously polls
// the connection for server
// responses. responses are then
// filled into the appropriate
// waiter's input buffer. it returns
// on the first error returned by Read()
func (c *Client) readLoop() {
	var seq uint64
	var sz int
	var frame fType
	var lead [leadSize]byte
	bwr := fwd.NewReaderSize(c.conn, 4096)

	for {
		if !c.do(bwr.ReadFull(lead[:])) {
			return
		}

		seq, frame, sz = readFrame(lead)

		// only accept fCMD and fRES frames;
		// they are routed to waiters
		// precisely the same way
		if frame != fCMD && frame != fRES {
			// ignore
			if !c.do(bwr.Skip(sz)) {
				return
			}
			continue
		}

		w := c.pending.remove(seq)
		if w == nil {
			if !c.do(bwr.Skip(sz)) {
				return
			}
			continue
		}

		// fill the waiters input
		// buffer and then notify
		if cap(w.in) >= sz {
			w.in = w.in[:sz]
		} else {
			w.in = make([]byte, sz)
		}

		if !c.do(bwr.ReadFull(w.in)) {
			return
		}

		// wakeup waiter w/
		// error from last
		// read call (usually nil)
		w.err = nil
		w.done.Unlock()
	}
}

func (c *Client) writeLoop() {
	bwr := fwd.NewWriterSize(c.conn, 4096)

	// the idea here is to
	// take advantage of buffering
	// small writes, but without
	// having to periodically
	// flush the buffer from
	// a separate goroutine.
	// instead, we write from
	// the queue as many times
	// as we can without blocking,
	// and then flush.
	for {
		// wait for one write
		wt, ok := <-c.writing
		if !ok {
			return
		}
		// write it
		if !c.do(bwr.Write(wt.in)) {
			return
		}
		// try to match this write
		// with other writes
	more:
		select {
		case another, ok := <-c.writing:
			if ok {
				if !c.do(bwr.Write(another.in)) {
					return
				}
				goto more
			} else {
				bwr.Flush()
				return
			}
		default:
			if !c.do(0, bwr.Flush()) {
				return
			}
		}
	}
}

// once every 'msec' milliseconds, reap
// every pending item with reap=true, and
// set all others to reap=true.
//
// NOTE(pmh): this doesn't actually guarantee
// de-queing after the given duration;
// rather, it limits max wait time to 2*msec.
func (c *Client) timeoutLoop(d time.Duration) {
	tick := time.Tick(d)
	for {
		select {
		case <-c.done:
			return
		case <-tick:
			c.pending.reap()
		}
	}
}

// write a command to the connection - works
// similarly to standard write()
func (w *waiter) writeCommand(cmd command, msg []byte) error {
	w.parent.wg.Add(1)
	if atomic.LoadUint32(&w.parent.state) == clientClosed {
		return ErrClosed
	}
	seqn := atomic.AddUint64(&w.parent.csn, 1)

	cmdlen := len(msg) + 1
	if cmdlen > maxMessageSize {
		return ErrTooLarge
	}

	// write frame + message
	need := leadSize + cmdlen
	if cap(w.in) >= need {
		w.in = w.in[:need]
	} else {
		w.in = make([]byte, need)
	}
	putFrame(w.in, seqn, fCMD, cmdlen)
	w.in[leadSize] = byte(cmd)
	copy(w.in[leadSize+1:], msg)

	w.reap = false
	w.seq = seqn

	p := w.parent
	p.pending.insert(w)
	p.writing <- w
	return nil
}

func (w *waiter) write(method string, in msgp.Marshaler) error {
	w.parent.wg.Add(1)
	if atomic.LoadUint32(&w.parent.state) == clientClosed {
		return ErrClosed
	}
	sn := atomic.AddUint64(&w.parent.csn, 1)

	var err error

	// save bytes up front
	if cap(w.in) < leadSize {
		w.in = make([]byte, leadSize, 256)
	} else {
		w.in = w.in[:leadSize]
	}

	// write body
	w.in = msgp.AppendString(w.in, method)
	// handle nil body
	if in != nil {
		w.in, err = in.MarshalMsg(w.in)
		if err != nil {
			return err
		}
	} else {
		w.in = msgp.AppendMapHeader(w.in, 0)
	}

	// raw request body
	olen := len(w.in) - leadSize

	if olen > maxMessageSize {
		return ErrTooLarge
	}

	putFrame(w.in, sn, fREQ, olen)
	w.reap = false
	w.seq = sn

	p := w.parent
	p.pending.insert(w)
	p.writing <- w
	return nil
}

func (w *waiter) read(out msgp.Unmarshaler) error {
	code, body, err := msgp.ReadIntBytes(w.in)
	if err != nil {
		return err
	}
	if Status(code) != okStatus {
		return Status(code)
	}
	if out != nil {
		_, err = out.UnmarshalMsg(body)
	}
	return err
}

func (w *waiter) call(method string, in msgp.Marshaler, out msgp.Unmarshaler) error {
	err := w.write(method, in)
	if err != nil {
		w.parent.wg.Done()
		return err
	}
	// wait for response
	w.done.Lock()
	w.parent.wg.Done()
	if w.err != nil {
		return w.err
	}
	return w.read(out)
}

// Call sends a request to the server with 'in' as the body,
// and then decodes the response into 'out'. Call is safe
// to call from multiple goroutines simultaneously.
func (c *Client) Call(method string, in msgp.Marshaler, out msgp.Unmarshaler) error {
	w := waiters.pop(c)
	err := w.call(method, in, out)
	waiters.push(w)
	return err
}

// errors specific to commands
var (
	errNoCmd      = errors.New("synapse: no response CMD code")
	errInvalidCmd = errors.New("synapse: invalid command")
	errUnknownCmd = errors.New("synapse: unknown command")
)

// doCommand executes a command from the client
func (c *Client) sendCommand(cmd command, msg []byte) error {
	w := waiters.pop(c)
	err := w.writeCommand(cmd, msg)
	if err != nil {
		c.wg.Done()
		return err
	}

	// wait
	w.done.Lock()
	c.wg.Done()
	if w.err != nil {
		return w.err
	}

	// bad response
	if len(w.in) == 0 {
		waiters.push(w)
		return errNoCmd
	}

	ret := command(w.in[0])

	if ret == cmdInvalid || ret >= _maxcommand {
		waiters.push(w)
		return errInvalidCmd
	}

	act := cmdDirectory[ret]
	if act == nil {
		waiters.push(w)
		return errUnknownCmd
	}

	act.Client(c, w.in[1:])
	waiters.push(w)
	return nil
}

// perform the ping command;
// returns an error if the server
// didn't respond appropriately
func (c *Client) ping() error {
	return c.sendCommand(cmdPing, nil)
}
<commit_msg>defaultTimeout is 3 seconds<commit_after>package synapse

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/philhofer/fwd"
	"github.com/tinylib/msgp/msgp"
)

const (
	defaultTimeout = 3 * time.Second // three seconds

	// waiter "high water mark"
	// TODO(maybe): make this adjustable.
	waiterHWM = 32
)

const (
	clientClosed = iota
	clientOpen
)

var (
	// ErrClosed is returns when a call is attempted
	// on a closed client.
	ErrClosed = errors.New("synapse: client is closed")

	// ErrTimeout is returned when a server
	// doesn't respond to a request before
	// the client's timeout scavenger can
	// free the waiting goroutine
	ErrTimeout = errors.New("synapse: the server didn't respond in time")

	// ErrTooLarge is returned when the message
	// size is larger than 65,535 bytes.
	ErrTooLarge = errors.New("synapse: message body too large")
)

// Dial creates a new client by dialing
// the provided network and remote address.
// The provided timeout is used as the timeout
// for requests, in milliseconds.
func Dial(network string, raddr string, timeout time.Duration) (*Client, error) {
	conn, err := net.Dial(network, raddr)
	if err != nil {
		return nil, err
	}
	return NewClient(conn, timeout)
}

// DialTLS acts identically to Dial, except that it dials the connection
// over TLS using the provided *tls.Config.
func DialTLS(network, raddr string, timeout time.Duration, config *tls.Config) (*Client, error) {
	conn, err := tls.Dial(network, raddr, config)
	if err != nil {
		return nil, err
	}
	return NewClient(conn, timeout)
}

// NewClient creates a new client from an
// existing net.Conn. Timeout is the maximum time,
// in milliseconds, to wait for server responses
// before sending an error to the caller. NewClient
// fails with an error if it cannot ping the server
// over the connection.
func NewClient(c net.Conn, timeout time.Duration) (*Client, error) {
	cl := &Client{
		conn:    c,
		writing: make(chan *waiter, waiterHWM),
		done:    make(chan struct{}),
		state:   clientOpen,
	}
	go cl.readLoop()
	go cl.writeLoop()

	if timeout <= 0 {
		timeout = defaultTimeout
	}

	go cl.timeoutLoop(timeout)

	// do a ping to check
	// for sanity
	err := cl.ping()
	if err != nil {
		cl.Close()
		return nil, fmt.Errorf("synapse: ping failed: %s", err)
	}

	return cl, nil
}

// Client is a client to
// a single synapse server.
type Client struct {
	conn    net.Conn       // connection
	csn     uint64         // sequence number; atomic
	writing chan *waiter   // queue to write to conn; size is effectively HWM
	done    chan struct{}  // closed during (*Client).Close to shut down timeoutloop
	wg      sync.WaitGroup // outstanding client procs
	state   uint32         // open, closed, etc.
	pending wMap           // map seq number to waiting handler
}

// used to transfer control
// flow to blocking goroutines
type waiter struct {
	next   *waiter    // next in linked list, or nil
	parent *Client    // parent *client
	seq    uint64     // sequence number
	done   sync.Mutex // for notifying response; locked is default
	err    error      // response error on wakeup, if applicable
	in     []byte     // response body
	reap   bool       // can reap for timeout
	static bool       // is part of the statically allocated arena
}

// Close idempotently closes the
// client's connection to the server.
// Close returns once every waiting
// goroutine has either received a
// response or timed out; goroutines
// with requests in progress are not interrupted.
func (c *Client) Close() error {
	// there can only be one winner of the Close() race
	if !atomic.CompareAndSwapUint32(&c.state, clientOpen, clientClosed) {
		return ErrClosed
	}

	c.wg.Wait()
	close(c.done)
	close(c.writing)
	return c.conn.Close()
}

// close with error
// sets the status of every waiting
// goroutine to 'err' and unblocks it.
func (c *Client) closeError(err error) {
	if !atomic.CompareAndSwapUint32(&c.state, clientOpen, clientClosed) {
		return
	}

	err = fmt.Errorf("synapse: fatal error: %s", err)

	// we can't actually guarantee that we will preempt
	// every goroutine, but we can try.
	c.pending.flush(err)
	for c.pending.length() > 0 {
		c.pending.flush(err)
	}

	c.wg.Wait()
	close(c.done)
	close(c.writing)
	c.conn.Close()
}

// a handler for io.Read() and io.Write(),
// e.g.
//  if !c.do(w.Write(data)) { goto fail }
func (c *Client) do(_ int, err error) bool {
	if err != nil {
		c.closeError(err)
		return false
	}
	return true
}

// readLoop continuously polls
// the connection for server
// responses. responses are then
// filled into the appropriate
// waiter's input buffer. it returns
// on the first error returned by Read()
func (c *Client) readLoop() {
	var seq uint64
	var sz int
	var frame fType
	var lead [leadSize]byte
	bwr := fwd.NewReaderSize(c.conn, 4096)

	for {
		if !c.do(bwr.ReadFull(lead[:])) {
			return
		}

		seq, frame, sz = readFrame(lead)

		// only accept fCMD and fRES frames;
		// they are routed to waiters
		// precisely the same way
		if frame != fCMD && frame != fRES {
			// ignore
			if !c.do(bwr.Skip(sz)) {
				return
			}
			continue
		}

		w := c.pending.remove(seq)
		if w == nil {
			if !c.do(bwr.Skip(sz)) {
				return
			}
			continue
		}

		// fill the waiters input
		// buffer and then notify
		if cap(w.in) >= sz {
			w.in = w.in[:sz]
		} else {
			w.in = make([]byte, sz)
		}

		if !c.do(bwr.ReadFull(w.in)) {
			return
		}

		// wakeup waiter w/
		// error from last
		// read call (usually nil)
		w.err = nil
		w.done.Unlock()
	}
}

func (c *Client) writeLoop() {
	bwr := fwd.NewWriterSize(c.conn, 4096)

	// the idea here is to
	// take advantage of buffering
	// small writes, but without
	// having to periodically
	// flush the buffer from
	// a separate goroutine.
	// instead, we write from
	// the queue as many times
	// as we can without blocking,
	// and then flush.
	for {
		// wait for one write
		wt, ok := <-c.writing
		if !ok {
			return
		}
		// write it
		if !c.do(bwr.Write(wt.in)) {
			return
		}
		// try to match this write
		// with other writes
	more:
		select {
		case another, ok := <-c.writing:
			if ok {
				if !c.do(bwr.Write(another.in)) {
					return
				}
				goto more
			} else {
				bwr.Flush()
				return
			}
		default:
			if !c.do(0, bwr.Flush()) {
				return
			}
		}
	}
}

// once every 'msec' milliseconds, reap
// every pending item with reap=true, and
// set all others to reap=true.
//
// NOTE(pmh): this doesn't actually guarantee
// de-queing after the given duration;
// rather, it limits max wait time to 2*msec.
func (c *Client) timeoutLoop(d time.Duration) {
	tick := time.Tick(d)
	for {
		select {
		case <-c.done:
			return
		case <-tick:
			c.pending.reap()
		}
	}
}

// write a command to the connection - works
// similarly to standard write()
func (w *waiter) writeCommand(cmd command, msg []byte) error {
	w.parent.wg.Add(1)
	if atomic.LoadUint32(&w.parent.state) == clientClosed {
		return ErrClosed
	}
	seqn := atomic.AddUint64(&w.parent.csn, 1)

	cmdlen := len(msg) + 1
	if cmdlen > maxMessageSize {
		return ErrTooLarge
	}

	// write frame + message
	need := leadSize + cmdlen
	if cap(w.in) >= need {
		w.in = w.in[:need]
	} else {
		w.in = make([]byte, need)
	}
	putFrame(w.in, seqn, fCMD, cmdlen)
	w.in[leadSize] = byte(cmd)
	copy(w.in[leadSize+1:], msg)

	w.reap = false
	w.seq = seqn

	p := w.parent
	p.pending.insert(w)
	p.writing <- w
	return nil
}

func (w *waiter) write(method string, in msgp.Marshaler) error {
	w.parent.wg.Add(1)
	if atomic.LoadUint32(&w.parent.state) == clientClosed {
		return ErrClosed
	}
	sn := atomic.AddUint64(&w.parent.csn, 1)

	var err error

	// save bytes up front
	if cap(w.in) < leadSize {
		w.in = make([]byte, leadSize, 256)
	} else {
		w.in = w.in[:leadSize]
	}

	// write body
	w.in = msgp.AppendString(w.in, method)
	// handle nil body
	if in != nil {
		w.in, err = in.MarshalMsg(w.in)
		if err != nil {
			return err
		}
	} else {
		w.in = msgp.AppendMapHeader(w.in, 0)
	}

	// raw request body
	olen := len(w.in) - leadSize

	if olen > maxMessageSize {
		return ErrTooLarge
	}

	putFrame(w.in, sn, fREQ, olen)
	w.reap = false
	w.seq = sn

	p := w.parent
	p.pending.insert(w)
	p.writing <- w
	return nil
}

func (w *waiter) read(out msgp.Unmarshaler) error {
	code, body, err := msgp.ReadIntBytes(w.in)
	if err != nil {
		return err
	}
	if Status(code) != okStatus {
		return Status(code)
	}
	if out != nil {
		_, err = out.UnmarshalMsg(body)
	}
	return err
}

func (w *waiter) call(method string, in msgp.Marshaler, out msgp.Unmarshaler) error {
	err := w.write(method, in)
	if err != nil {
		w.parent.wg.Done()
		return err
	}
	// wait for response
	w.done.Lock()
	w.parent.wg.Done()
	if w.err != nil {
		return w.err
	}
	return w.read(out)
}

// Call sends a request to the server with 'in' as the body,
// and then decodes the response into 'out'. Call is safe
// to call from multiple goroutines simultaneously.
func (c *Client) Call(method string, in msgp.Marshaler, out msgp.Unmarshaler) error {
	w := waiters.pop(c)
	err := w.call(method, in, out)
	waiters.push(w)
	return err
}

// errors specific to commands
var (
	errNoCmd      = errors.New("synapse: no response CMD code")
	errInvalidCmd = errors.New("synapse: invalid command")
	errUnknownCmd = errors.New("synapse: unknown command")
)

// doCommand executes a command from the client
func (c *Client) sendCommand(cmd command, msg []byte) error {
	w := waiters.pop(c)
	err := w.writeCommand(cmd, msg)
	if err != nil {
		c.wg.Done()
		return err
	}

	// wait
	w.done.Lock()
	c.wg.Done()
	if w.err != nil {
		return w.err
	}

	// bad response
	if len(w.in) == 0 {
		waiters.push(w)
		return errNoCmd
	}

	ret := command(w.in[0])

	if ret == cmdInvalid || ret >= _maxcommand {
		waiters.push(w)
		return errInvalidCmd
	}

	act := cmdDirectory[ret]
	if act == nil {
		waiters.push(w)
		return errUnknownCmd
	}

	act.Client(c, w.in[1:])
	waiters.push(w)
	return nil
}

// perform the ping command;
// returns an error if the server
// didn't respond appropriately
func (c *Client) ping() error {
	return c.sendCommand(cmdPing, nil)
}
<|endoftext|>


<commit_before>package titan

import (
	"sync"

	"github.com/neptulon/cmap"
	"github.com/neptulon/jsonrpc"
	"github.com/neptulon/neptulon"
)

// Client is a Titan API client.
type Client struct {
	client *jsonrpc.Client
}

// NewClient creates a new Client object.
// msgWG = (optional) sets the given *sync.WaitGroup reference to be used for counting active gorotuines that are used for handling incoming/outgoing messages.
// disconnHandler = (optional) registers a function to handle client disconnection events.
func NewClient(msgWG *sync.WaitGroup, disconnHandler func(client *neptulon.Client)) *Client {
	return &Client{client: jsonrpc.NewClient(msgWG, disconnHandler)}
}

// Connect connectes to the server at given network address and starts receiving messages.
func (c *Client) Connect(addr string, debug bool) error {
	return c.client.Connect(addr, debug)
}

// ConnID is a randomly generated unique client connection ID.
func (c *Client) ConnID() string {
	return c.client.ConnID()
}

// Session is a thread-safe data store for storing arbitrary data for this connection session.
func (c *Client) Session() *cmap.CMap {
	return c.client.Session()
}

// SetDeadline set the read/write deadlines for the connection, in seconds.
func (c *Client) SetDeadline(seconds int) {
	c.client.SetDeadline(seconds)
}

// Close closes a client connection.
func (c *Client) Close() error {
	return c.client.Close()
}
<commit_msg>add Client.ReceiveMessages prototype<commit_after>package titan

import (
	"sync"
	"time"

	"github.com/neptulon/cmap"
	"github.com/neptulon/jsonrpc"
	"github.com/neptulon/neptulon"
)

// Client is a Titan API client.
type Client struct {
	client *jsonrpc.Client
}

// NewClient creates a new Client object.
// msgWG = (optional) sets the given *sync.WaitGroup reference to be used for counting active gorotuines that are used for handling incoming/outgoing messages.
// disconnHandler = (optional) registers a function to handle client disconnection events.
func NewClient(msgWG *sync.WaitGroup, disconnHandler func(client *neptulon.Client)) *Client {
	return &Client{client: jsonrpc.NewClient(msgWG, disconnHandler)}
}

// Connect connectes to the server at given network address and starts receiving messages.
func (c *Client) Connect(addr string, debug bool) error {
	return c.client.Connect(addr, debug)
}

// ConnID is a randomly generated unique client connection ID.
func (c *Client) ConnID() string {
	return c.client.ConnID()
}

// Session is a thread-safe data store for storing arbitrary data for this connection session.
func (c *Client) Session() *cmap.CMap {
	return c.client.Session()
}

// SetDeadline set the read/write deadlines for the connection, in seconds.
func (c *Client) SetDeadline(seconds int) {
	c.client.SetDeadline(seconds)
}

// Close closes a client connection.
func (c *Client) Close() error {
	return c.client.Close()
}

// RecvMsg is an incoming chat message.
type RecvMsg struct {
	From    string    `json:"from"`
	Time    time.Time `json:"time"`
	Message string    `json:"message"`
}

// ReceiveMessages sends a request to server to receive any pending messages.
func (c *Client) ReceiveMessages(msgHandler func(m []RecvMsg) error) error {
	_, err := c.client.SendRequest("msg.recv", nil, func(ctx *jsonrpc.ResCtx) error {
		var msg []RecvMsg
		if err := ctx.Result(msg); err != nil {
			return err
		}

		if err := msgHandler(msg); err != nil {
			return err
		}

		return ctx.Next()
	})

	return err
}
<|endoftext|>


<commit_before>package consultant

import (
	"errors"
	"fmt"
	"github.com/hashicorp/consul/api"
	"github.com/renstrom/shortuuid"
	"math/rand"
	"net/url"
	"os"
	"strings"
)

type Client struct {
	*api.Client

	myaddr string
	myhost string

	logSlug      string
	logSlugSlice []interface{}
}

// NewConnection generates a fresh connection to the local consul agent using the default config
func NewClient(conf *api.Config) (*Client, error) {
	var err error

	client := &Client{
		logSlug:      "[consultant-client]",
		logSlugSlice: []interface{}{"[consultant-client]"},
	}

	client.Client, err = api.NewClient(conf)
	if err != nil {
		return nil, fmt.Errorf("Unable to create Consul API Client: %v", err)
	}

	if client.myhost, err = os.Hostname(); err != nil {
		client.logPrintf("Unable to determine hostname: %v", err)
	}

	if client.myaddr, err = GetMyAddress(); err != nil {
		client.logPrintf("Unable to determine ip address: %v", err)
	}

	return client, nil
}

// Service will attempt to locate any registered service with a name + tag combination and return one at random from
// the resulting list
func (c *Client) Service(service, tag string, passingOnly bool, options *api.QueryOptions) (*api.ServiceEntry, *api.QueryMeta, error) {
	svcs, qm, err := c.Health().Service(service, tag, passingOnly, options)
	if nil != err {
		return nil, nil, err
	}

	svcLen := len(svcs)
	if 0 < svcLen {
		return svcs[rand.Intn(svcLen)], qm, nil
	}

	return nil, qm, nil
}

// SimpleServiceRegistration describes a service that we want to register
type SimpleServiceRegistration struct {
	Name string // [required] name to register service under
	Port int    // [required] external port to advertise for service consumers
	Id   string // [optional] specific id for service, will be generated if not set

	RandomId          bool     // [optional] if Id is not set, use a random uuid if true, or hostname if false
	Address           string   // [optional] determined automatically by Register() if not set
	Tags              []string // [optional] desired tags: Register() adds serviceId
	CheckTCP          bool     // [optional] if true register a TCP check
	CheckPath         string   // [optional] register an http check with this path if set
	CheckScheme       string   // [optional] override the http check scheme (default: http)
	CheckPort         int      // [optional] if set, this is the port that the health check lives at
	Interval          string   // [optional] check interval
	EnableTagOverride bool     // [optional] whether we should allow tag overriding (new in 0.6+)
}

// SimpleServiceRegister is a helper method to ease consul service registration
func (c *Client) SimpleServiceRegister(reg *SimpleServiceRegistration) (string, error) {
	var err error                        // generic error holder
	var serviceId string                 // local service identifier
	var address string                   // service host address
	var interval string                  // check interval
	var checkHTTP *api.AgentServiceCheck // http type check
	var checkTCP *api.AgentServiceCheck  // tcp port check
	var serviceName string               // service registration name

	// Perform some basic service name cleanup and validation
	serviceName = strings.TrimSpace(reg.Name)
	if serviceName == "" {
		return "", errors.New("\"Name\" cannot be blank")
	}
	if strings.Contains(serviceName, " ") {
		return "", fmt.Errorf("Specified service name \"%s\" is invalid, service names cannot contain spaces", serviceName)
	}

	// Come on, guys...valid ports plz...
	if reg.Port <= 0 {
		return "", fmt.Errorf("%d is not a valid port", reg.Port)
	}

	if address = reg.Address; address == "" {
		address = c.myaddr
	}

	if serviceId = reg.Id; serviceId == "" {
		// Form a unique service id
		var tail string
		if reg.RandomId {
			tail = shortuuid.New()
		} else {
			tail = strings.ToLower(c.myhost)
		}
		serviceId = fmt.Sprintf("%s-%s", serviceName, tail)
	}

	if interval = reg.Interval; interval == "" {
		// set a default interval
		interval = "30s"
	}

	// The serviceId is added in order to ensure detection in ServiceMonitor()
	tags := append(reg.Tags, serviceId)

	// Set up the service registration struct
	asr := &api.AgentServiceRegistration{
		ID:                serviceId,
		Name:              serviceName,
		Tags:              tags,
		Port:              reg.Port,
		Address:           address,
		Checks:            api.AgentServiceChecks{},
		EnableTagOverride: reg.EnableTagOverride,
	}

	// allow port override
	checkPort := reg.CheckPort
	if checkPort <= 0 {
		checkPort = reg.Port
	}

	// build http check if specified
	if reg.CheckPath != "" {

		// allow scheme override
		checkScheme := reg.CheckScheme
		if checkScheme == "" {
			checkScheme = "http"
		}

		// build check url
		checkURL := &url.URL{
			Scheme: checkScheme,
			Host:   fmt.Sprintf("%s:%d", address, checkPort),
			Path:   reg.CheckPath,
		}

		// build check
		checkHTTP = &api.AgentServiceCheck{
			HTTP:     checkURL.String(),
			Interval: reg.Interval,
		}

		// add http check
		asr.Checks = append(asr.Checks, checkHTTP)
	}

	// build tcp check if specified
	if reg.CheckTCP {
		// create tcp check definition
		checkTCP = &api.AgentServiceCheck{
			TCP:      fmt.Sprintf("%s:%d", address, checkPort),
			Interval: interval,
		}
	}

	// add tcp check if defined
	if checkTCP != nil {
		asr.Checks = append(asr.Checks, checkTCP)
	}

	// register and check error
	if err = c.Agent().ServiceRegister(asr); err != nil {
		return "", err
	}

	// return registered service
	return serviceId, nil
}

func (c *Client) logPrintf(format string, v ...interface{}) {
	log.Printf(fmt.Sprintf("%s %s", c.logSlug, format), v...)
}

func (c *Client) logPrint(v ...interface{}) {
	log.Print(append(c.logSlugSlice, v...)...)
}

func (c *Client) logPrintln(v ...interface{}) {
	log.Println(append(c.logSlugSlice, v...)...)
}

func (c *Client) logFatalf(format string, v ...interface{}) {
	log.Fatalf(fmt.Sprintf("%s %s", c.logSlug, format), v...)
}

func (c *Client) logFatal(v ...interface{}) {
	log.Fatal(append(c.logSlugSlice, v...)...)
}

func (c *Client) logFatalln(v ...interface{}) {
	log.Fatalln(append(c.logSlugSlice, v...)...)
}

func (c *Client) logPanicf(format string, v ...interface{}) {
	log.Panicf(fmt.Sprintf("%s %s", c.logSlug, format), v...)
}

func (c *Client) logPanic(v ...interface{}) {
	log.Panic(append(c.logSlugSlice, v...)...)
}

func (c *Client) logPanicln(v ...interface{}) {
	log.Panicln(append(c.logSlugSlice, v...)...)
}
<commit_msg>Adding client Addr and Host getter/setter pairs<commit_after>package consultant

import (
	"errors"
	"fmt"
	"github.com/hashicorp/consul/api"
	"github.com/renstrom/shortuuid"
	"math/rand"
	"net/url"
	"os"
	"strings"
)

type Client struct {
	*api.Client

	myAddr string
	myHost string

	logSlug      string
	logSlugSlice []interface{}
}

// NewConnection generates a fresh connection to the local consul agent using the default config
func NewClient(conf *api.Config) (*Client, error) {
	var err error

	client := &Client{
		logSlug:      "[consultant-client]",
		logSlugSlice: []interface{}{"[consultant-client]"},
	}

	client.Client, err = api.NewClient(conf)
	if err != nil {
		return nil, fmt.Errorf("Unable to create Consul API Client: %v", err)
	}

	if client.myHost, err = os.Hostname(); err != nil {
		client.logPrintf("Unable to determine hostname: %v", err)
	}

	if client.myAddr, err = GetMyAddress(); err != nil {
		client.logPrintf("Unable to determine ip address: %v", err)
	}

	return client, nil
}

func (c *Client) MyAddr() string {
	return c.myAddr
}

func (c *Client) SetMyAddr(myAddr string) {
	c.myAddr = myAddr
}

func (c *Client) MyHost() string {
	return c.myHost
}

func (c *Client) SetMyHost(myHost string) {
	c.myHost = myHost
}

// Service will attempt to locate any registered service with a name + tag combination and return one at random from
// the resulting list
func (c *Client) Service(service, tag string, passingOnly bool, options *api.QueryOptions) (*api.ServiceEntry, *api.QueryMeta, error) {
	svcs, qm, err := c.Health().Service(service, tag, passingOnly, options)
	if nil != err {
		return nil, nil, err
	}

	svcLen := len(svcs)
	if 0 < svcLen {
		return svcs[rand.Intn(svcLen)], qm, nil
	}

	return nil, qm, nil
}

// SimpleServiceRegistration describes a service that we want to register
type SimpleServiceRegistration struct {
	Name string // [required] name to register service under
	Port int    // [required] external port to advertise for service consumers
	Id   string // [optional] specific id for service, will be generated if not set

	RandomId          bool     // [optional] if Id is not set, use a random uuid if true, or hostname if false
	Address           string   // [optional] determined automatically by Register() if not set
	Tags              []string // [optional] desired tags: Register() adds serviceId
	CheckTCP          bool     // [optional] if true register a TCP check
	CheckPath         string   // [optional] register an http check with this path if set
	CheckScheme       string   // [optional] override the http check scheme (default: http)
	CheckPort         int      // [optional] if set, this is the port that the health check lives at
	Interval          string   // [optional] check interval
	EnableTagOverride bool     // [optional] whether we should allow tag overriding (new in 0.6+)
}

// SimpleServiceRegister is a helper method to ease consul service registration
func (c *Client) SimpleServiceRegister(reg *SimpleServiceRegistration) (string, error) {
	var err error                        // generic error holder
	var serviceId string                 // local service identifier
	var address string                   // service host address
	var interval string                  // check interval
	var checkHTTP *api.AgentServiceCheck // http type check
	var checkTCP *api.AgentServiceCheck  // tcp port check
	var serviceName string               // service registration name

	// Perform some basic service name cleanup and validation
	serviceName = strings.TrimSpace(reg.Name)
	if serviceName == "" {
		return "", errors.New("\"Name\" cannot be blank")
	}
	if strings.Contains(serviceName, " ") {
		return "", fmt.Errorf("Specified service name \"%s\" is invalid, service names cannot contain spaces", serviceName)
	}

	// Come on, guys...valid ports plz...
	if reg.Port <= 0 {
		return "", fmt.Errorf("%d is not a valid port", reg.Port)
	}

	if address = reg.Address; address == "" {
		address = c.myAddr
	}

	if serviceId = reg.Id; serviceId == "" {
		// Form a unique service id
		var tail string
		if reg.RandomId {
			tail = shortuuid.New()
		} else {
			tail = strings.ToLower(c.myHost)
		}
		serviceId = fmt.Sprintf("%s-%s", serviceName, tail)
	}

	if interval = reg.Interval; interval == "" {
		// set a default interval
		interval = "30s"
	}

	// The serviceId is added in order to ensure detection in ServiceMonitor()
	tags := append(reg.Tags, serviceId)

	// Set up the service registration struct
	asr := &api.AgentServiceRegistration{
		ID:                serviceId,
		Name:              serviceName,
		Tags:              tags,
		Port:              reg.Port,
		Address:           address,
		Checks:            api.AgentServiceChecks{},
		EnableTagOverride: reg.EnableTagOverride,
	}

	// allow port override
	checkPort := reg.CheckPort
	if checkPort <= 0 {
		checkPort = reg.Port
	}

	// build http check if specified
	if reg.CheckPath != "" {

		// allow scheme override
		checkScheme := reg.CheckScheme
		if checkScheme == "" {
			checkScheme = "http"
		}

		// build check url
		checkURL := &url.URL{
			Scheme: checkScheme,
			Host:   fmt.Sprintf("%s:%d", address, checkPort),
			Path:   reg.CheckPath,
		}

		// build check
		checkHTTP = &api.AgentServiceCheck{
			HTTP:     checkURL.String(),
			Interval: reg.Interval,
		}

		// add http check
		asr.Checks = append(asr.Checks, checkHTTP)
	}

	// build tcp check if specified
	if reg.CheckTCP {
		// create tcp check definition
		checkTCP = &api.AgentServiceCheck{
			TCP:      fmt.Sprintf("%s:%d", address, checkPort),
			Interval: interval,
		}
	}

	// add tcp check if defined
	if checkTCP != nil {
		asr.Checks = append(asr.Checks, checkTCP)
	}

	// register and check error
	if err = c.Agent().ServiceRegister(asr); err != nil {
		return "", err
	}

	// return registered service
	return serviceId, nil
}

func (c *Client) logPrintf(format string, v ...interface{}) {
	log.Printf(fmt.Sprintf("%s %s", c.logSlug, format), v...)
}

func (c *Client) logPrint(v ...interface{}) {
	log.Print(append(c.logSlugSlice, v...)...)
}

func (c *Client) logPrintln(v ...interface{}) {
	log.Println(append(c.logSlugSlice, v...)...)
}

func (c *Client) logFatalf(format string, v ...interface{}) {
	log.Fatalf(fmt.Sprintf("%s %s", c.logSlug, format), v...)
}

func (c *Client) logFatal(v ...interface{}) {
	log.Fatal(append(c.logSlugSlice, v...)...)
}

func (c *Client) logFatalln(v ...interface{}) {
	log.Fatalln(append(c.logSlugSlice, v...)...)
}

func (c *Client) logPanicf(format string, v ...interface{}) {
	log.Panicf(fmt.Sprintf("%s %s", c.logSlug, format), v...)
}

func (c *Client) logPanic(v ...interface{}) {
	log.Panic(append(c.logSlugSlice, v...)...)
}

func (c *Client) logPanicln(v ...interface{}) {
	log.Panicln(append(c.logSlugSlice, v...)...)
}
<|endoftext|>


<commit_before>package wsevent

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

	ws "github.com/gorilla/websocket"
)

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

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

type request struct {
	Id   string
	Data json.RawMessage
}

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

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

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

	return base64.URLEncoding.EncodeToString(bytes)
}

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

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

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

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

	return client, nil
}

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

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

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

type emitJS struct {
	Id   int         `json:"id"`
	Data interface{} `json:"data"`
}

var emitPool = &sync.Pool{New: func() interface{} { return emitJS{} }}

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

	js := emitPool.Get().(emitJS)
	defer emitPool.Put(js)

	js.Id = -1
	js.Data = v

	return c.conn.WriteJSON(js)
}

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

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

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

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

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

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

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

		js := reqPool.Get().(request)

		if err := json.Unmarshal(data, &js); err != nil {
			continue
		}

		callName := s.Extractor(js.Data)

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

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

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

			reqPool.Put(js)
			replyPool.Put(replyJs)
		}()
		time.Sleep(10 * time.Millisecond)
	}
}
<commit_msg>Use channels for receiving and sending<commit_after>package wsevent

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

	ws "github.com/gorilla/websocket"
)

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

	conn    *ws.Conn
	send    chan []byte
	request *http.Request
}

type request struct {
	Id   string
	Data json.RawMessage
}

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

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

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

	return base64.URLEncoding.EncodeToString(bytes)
}

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

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

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

	client := &Client{
		id:      id,
		conn:    conn,
		send:    make(chan []byte),
		request: r,
	}
	s.newClient <- client

	return client, nil
}

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

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) {
	c.send <- []byte(data)
}

type emitJS struct {
	Id   int         `json:"id"`
	Data interface{} `json:"data"`
}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	js := emitJS{}
	js.Id = -1
	js.Data = v

	bytes, err := json.Marshal(js)
	if err != nil {
		return err
	}

	c.send <- bytes
	return nil
}

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

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

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

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

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

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

func (c *Client) listener(s *Server) {
	close := make(chan struct{})
	recv := make(chan []byte)
	tick := time.NewTicker(time.Millisecond * 10)
	go func() {
		for {
			select {
			case data := <-c.send:
				c.conn.WriteMessage(ws.TextMessage, data)
			case <-close:
				return
			}
		}
	}()
	go func() {
		for {
			select {
			case <-tick.C:
				mtype, data, err := c.conn.ReadMessage()
				if err != nil {
					c.cleanup(s)
					recv <- []byte{}
					close <- struct{}{}
					tick.Stop()
					return
				}
				if mtype != ws.TextMessage {
					continue
				}

				recv <- data
			}
		}
	}()

	for {
		data := <-recv
		if len(data) == 0 {
			return
		}

		req := reqPool.Get().(request)

		if err := json.Unmarshal(data, &req); err != nil {
			continue
		}

		callName := s.Extractor(req.Data)

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

		if !ok {
			if s.DefaultHandler != nil {
				f = s.DefaultHandler
				goto call
			}
			continue
		}
	call:
		go func() {
			rtrn := f(s, c, req.Data)
			reply := replyPool.Get().(reply)
			reply.Id = req.Id
			reply.Data = rtrn

			bytes, _ := json.Marshal(reply)
			c.send <- bytes

			reqPool.Put(req)
			replyPool.Put(reply)
		}()
	}
}
<|endoftext|>


<commit_before>package main

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

type VenueStatus struct {
	OK    bool `json:"ok"`
	Venue string `json:"venue"`
}

type APIStatus struct {
	OK    bool `json:"ok"`
	Error string `json:"error"`
}

func main() {
	// first call to Stockfighter API
	endpoint := "https://api.stockfighter.io/ob/api/heartbeat"
	bodytext := stockfighterAPI(endpoint)
	fmt.Printf("Stockfighter heartbeat: %s\n", bodytext)

	var apiStates APIStatus
	err := json.Unmarshal(bodytext, &apiStates)
	if err != nil {
		panic(err)
	}
	fmt.Printf("API heartbeat: %+v\n", apiStates)

	venueHeartbeat := "https://api.stockfighter.io/ob/api/venues/TESTEX/heartbeat"
	bodytext = stockfighterAPI(venueHeartbeat)
	fmt.Printf("TESTEX heartbeat: %s\n", bodytext)

	var venueStates VenueStatus
	err = json.Unmarshal(bodytext, &venueStates)
	if err != nil {
		panic(err)
	}

	fmt.Printf("TESTEX heartbeat: %+v\n", venueStates)
}

func stockfighterAPI(endpoint string) []byte {
	res, err := http.Get(endpoint)

	if err != nil {
		panic(err)
	}
	bodytext, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		panic(err)
	}
	return bodytext
}
<commit_msg>Drive through to order book API call<commit_after>package main

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

type VenueStatus struct {
	OK    bool `json:"ok"`
	Venue string `json:"venue"`
}

type APIStatus struct {
	OK    bool `json:"ok"`
	Error string `json:"error"`
}

type StockDetails struct {
	Name    string `json:"name"`
	Symbol string `json:"symbol"`
}

type VenueStocks struct {
	OK    bool `json:"ok"`
	Stocks []StockDetails `json:"symbols"`
}

type Order struct {
	Price    int `json:"price"`
	Quantity int `json:"qty"`
	IsBuy    bool `json:"isBuy"`
}

type OrderBook struct {
	OK    bool `json:"ok"`
	Venue string `json:"venue"`
	Symbol string `json:"symbol"`
	Timestamp string `json:"ts"`

	Bids []Order `json:"bids"`
	Asks []Order `json:"asks"`
}
func main() {
	// first call to Stockfighter API
	endpoint := "https://api.stockfighter.io/ob/api/heartbeat"
	bodytext := stockfighterAPI(endpoint)
	fmt.Printf("Stockfighter heartbeat: %s\n", bodytext)

	var apiStates APIStatus
	err := json.Unmarshal(bodytext, &apiStates)
	if err != nil {
		panic(err)
	}
	fmt.Printf("API heartbeat: %+v\n", apiStates)

	venueHeartbeat := "https://api.stockfighter.io/ob/api/venues/TESTEX/heartbeat"
	bodytext = stockfighterAPI(venueHeartbeat)
	fmt.Printf("TESTEX heartbeat: %s\n", bodytext)

	var venueStates VenueStatus
	err = json.Unmarshal(bodytext, &venueStates)
	if err != nil {
		panic(err)
	}

	fmt.Printf("TESTEX heartbeat: %+v\n", venueStates)

	venueStock := "https://api.stockfighter.io/ob/api/venues/TESTEX/stocks"
	bodytext = stockfighterAPI(venueStock)
	fmt.Printf("TESTEX stocks: %s\n", bodytext)

	var venueStocks VenueStocks
	err = json.Unmarshal(bodytext, &venueStocks)
	if err != nil {
		panic(err)
	}

	fmt.Printf("TESTEX stocks: %+v\n", venueStocks)

	stock := venueStocks.Stocks[0].Symbol
	stockOrder := "https://api.stockfighter.io/ob/api/venues/TESTEX/stocks/" + stock
	bodytext = stockfighterAPI(stockOrder)
	fmt.Printf("TESTEX order book: %s\n", bodytext)

	var orderbook OrderBook
	err = json.Unmarshal(bodytext, &orderbook)
	if err != nil {
		panic(err)
	}

	fmt.Printf("TESTEX order book: %+v\n", orderbook)

}

func stockfighterAPI(endpoint string) []byte {
	res, err := http.Get(endpoint)

	if err != nil {
		panic(err)
	}
	bodytext, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		panic(err)
	}
	return bodytext
}
<|endoftext|>


<commit_before>package turnpike

import (
	"fmt"
	"time"
)

const (
	PUBLISHER = 1 << iota
	SUBSCRIBER
	CALLEE
	CALLER
	ALL = PUBLISHER | SUBSCRIBER | CALLEE | CALLER
)

type Client struct {
	Peer
	ReceiveTimeout time.Duration
	roles          int
	listeners      map[ID]chan Message
	events         map[ID]EventHandler
	requestCount   uint
}

func NewWebsocketClient(serialization int, url string, realm URI, roles int) (*Client, error) {
	p, err := NewWebsocketPeer(serialization, url, "")
	if err != nil {
		return nil, err
	}
	c := &Client{
		Peer:           p,
		ReceiveTimeout: 5 * time.Second,
		roles:          roles,
		listeners:      make(map[ID]chan Message),
		events:         make(map[ID]EventHandler),
		requestCount:   0,
	}
	go c.Receive()
	// TODO: add roles to hello
	c.Send(&Hello{Realm: realm})
	return c, nil
}

func (c *Client) nextID() ID {
	c.requestCount++
	return ID(c.requestCount)
}

func (c *Client) Receive() {
	for msg := range c.Peer.Receive() {
		switch msg := msg.(type) {
		case *Event:
			if fn, ok := c.events[msg.Subscription]; ok {
				go fn(msg.Arguments, msg.ArgumentsKw)
			} else {
				log.Println("no handler registered for subscription:", msg.Subscription)
			}
		case *Subscribed:
			if l, ok := c.listeners[msg.Request]; ok {
				l <- msg
			} else {
				log.Println("no listener for subscribed:", msg.Request)
			}
		default:
			log.Println(msg.MessageType(), msg)
		}
	}
	log.Fatal("Receive buffer closed")
}

func (c *Client) registerListener(id ID) {
	log.Println("register listener:", id)
	wait := make(chan Message, 1)
	c.listeners[id] = wait
}

func (c *Client) waitOnListener(id ID) (msg Message, err error) {
	log.Println("wait on listener:", id)
	if wait, ok := c.listeners[id]; !ok {
		return nil, fmt.Errorf("unknown listener ID: %v", id)
	} else {
		select {
		case msg = <-wait:
			return
		case <-time.After(c.ReceiveTimeout):
			err = fmt.Errorf("timeout while waiting for message")
			return
		}
	}
}

type EventHandler func(args []interface{}, kwargs map[string]interface{})

func (c *Client) Subscribe(topic URI, fn EventHandler) error {
	id := c.nextID()
	c.registerListener(id)
	if err := c.Send(&Subscribe{Request: id, Topic: topic}); err != nil {
		return err
	}
	// wait to receive SUBSCRIBED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error subscribing to topic '%v': %v", topic, e.Error)
	} else if subscribed, ok := msg.(*Subscribed); !ok {
		return fmt.Errorf("unexpected message received: %s", msg.MessageType())
	} else {
		// register the event handler with this subscription
		c.events[subscribed.Subscription] = fn
	}
	return nil
}

func (c *Client) Publish(topic URI, args []interface{}, kwargs map[string]interface{}) error {
	return c.Send(&Publish{
		Request:     c.nextID(),
		Topic:       topic,
		Arguments:   args,
		ArgumentsKw: kwargs,
	})
}
<commit_msg>Begin work on adding RPC support:   - Add the Call method for invoking a remote RPC   - Have the HELLO message include the roles   - Add a WaitForSession method which waits for the session WELCOME message. Callers can use this to make sure a valid session exists before starting RPC calls.<commit_after>package turnpike

import (
	"fmt"
	"time"
)

const (
	PUBLISHER = 1 << iota
	SUBSCRIBER
	CALLEE
	CALLER
	ALL = PUBLISHER | SUBSCRIBER | CALLEE | CALLER
)

type Client struct {
	Peer
	ReceiveTimeout time.Duration
	roles          int
	listeners      map[ID]chan Message
	events         map[ID]EventHandler
	calls          map[ID]chan Message
	welcome        chan Message
	requestCount   uint
}

func NewWebsocketClient(serialization int, url string, realm URI, roles int) (*Client, error) {
	p, err := NewWebsocketPeer(serialization, url, "")
	if err != nil {
		return nil, err
	}
	c := &Client{
		Peer:           p,
		ReceiveTimeout: 5 * time.Second,
		roles:          roles,
		listeners:      make(map[ID]chan Message),
		calls:          make(map[ID]chan Message),
		events:         make(map[ID]EventHandler),
		welcome:        make(chan Message),
		requestCount:   0,
	}
	go c.Receive()

	roles_map := make(map[string]interface{})
	roles_map["publisher"] = make(map[string]interface{})
	roles_map["subscriber"] = make(map[string]interface{})
	roles_map["callee"] = make(map[string]interface{})
	roles_map["caller"] = make(map[string]interface{})

	details := make(map[string]interface{})
	details["roles"] = roles_map

	c.Send(&Hello{Realm: realm, Details: details})
	return c, nil
}

func (c *Client) WaitForSession() (msg Message, err error) {
	// wait to receive WELCOME message
	select {
	case msg := <-c.welcome:
		return msg, nil
	case <-time.After(c.ReceiveTimeout):
		return nil, fmt.Errorf("timeout while waiting for message")
	}
}

func (c *Client) nextID() ID {
	c.requestCount++
	return ID(c.requestCount)
}

func (c *Client) Receive() {
	for msg := range c.Peer.Receive() {
		switch msg := msg.(type) {
		case *Event:
			if fn, ok := c.events[msg.Subscription]; ok {
				go fn(msg.Arguments, msg.ArgumentsKw)
			} else {
				log.Println("no handler registered for subscription:", msg.Subscription)
			}

		case *Result:
			if l, ok := c.calls[msg.Request]; ok {
				l <- msg
			} else {
				log.Println("no handler registered for call:", msg.Request)
			}

		case *Welcome:
			c.welcome <- msg

		case *Subscribed:
			if l, ok := c.listeners[msg.Request]; ok {
				l <- msg
			} else {
				log.Println("no listener for subscribed:", msg.Request)
			}
		default:
			log.Println(msg.MessageType(), msg)
		}
	}
	log.Fatal("Receive buffer closed")
}

func (c *Client) registerListener(id ID) {
	log.Println("register listener:", id)
	wait := make(chan Message, 1)
	c.listeners[id] = wait
}

func (c *Client) waitOnListener(id ID) (msg Message, err error) {
	log.Println("wait on listener:", id)
	if wait, ok := c.listeners[id]; !ok {
		return nil, fmt.Errorf("unknown listener ID: %v", id)
	} else {
		select {
		case msg = <-wait:
			return
		case <-time.After(c.ReceiveTimeout):
			err = fmt.Errorf("timeout while waiting for message")
			return
		}
	}
}

type EventHandler func(args []interface{}, kwargs map[string]interface{})

func (c *Client) Subscribe(topic URI, fn EventHandler) error {
	id := c.nextID()
	c.registerListener(id)
	if err := c.Send(&Subscribe{Request: id, Topic: topic}); err != nil {
		return err
	}
	// wait to receive SUBSCRIBED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error subscribing to topic '%v': %v", topic, e.Error)
	} else if subscribed, ok := msg.(*Subscribed); !ok {
		return fmt.Errorf("unexpected message received: %s", msg.MessageType())
	} else {
		// register the event handler with this subscription
		c.events[subscribed.Subscription] = fn
	}
	return nil
}

func (c *Client) Publish(topic URI, args []interface{}, kwargs map[string]interface{}) error {
	return c.Send(&Publish{
		Request:     c.nextID(),
		Topic:       topic,
		Arguments:   args,
		ArgumentsKw: kwargs,
	})
}

func (c *Client) Call(procedure URI, args []interface{}, kwargs map[string]interface{}) (msg Message, err error) {
	callId := c.nextID()
	result := make(chan Message, 1)

	c.calls[callId] = result

	if err := c.Send(&Call{
		Request:     callId,
		Procedure:   procedure,
		Options:     make(map[string]interface{}),
		Arguments:   args,
		ArgumentsKw: kwargs,
	}); err != nil {
		return nil, err
	}

	// wait to receive RESULT message
	select {
	case msg = <-result:
		return msg, nil
	case <-time.After(c.ReceiveTimeout):
		return nil, fmt.Errorf("timeout while waiting for message")
	}
}
<|endoftext|>