<commit_before>/*
MAF Latency Package

Measure the latency of HTTP requests.

Usage

First specify a new http.Request like:

req, _ := http.NewRequest("GET", "http://google.com", nil)

To use this request with the latency library you have to embed it in a
LatencyRequest and specify a timeout interval. Example:

request := LatencyRequest{
	Request: req,
	Timeout: 10 * time.Second,
}

To execute the latency measurement, simple call the Execute() method:

response, err := request.Execute()

Note that the latency measurement follows redirects, but only measures the time
of the last non-redirecting request. The call will return a LatencyResponse
object.
*/
package latency

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"
)

var (
	ErrLatencyRedirect = "Redirect in request"
)

// LatencyRequest is a wrapper for http.Request.
// You can specify an additonal timeout
type LatencyRequest struct {
	*http.Request
	Timeout   time.Duration
	Redirects []string
}

func (l *LatencyRequest) String() string {
	return fmt.Sprintf("%v %v, timout is %v", l.Method, l.URL, l.Timeout)
}

// LatencyRedirectError is an error type used, if redirection occured during
// latency measurement
type LatencyRedirectError string

func (l LatencyRedirectError) Error() string {
	return ErrLatencyRedirect
}

// Redirect policy for http.Client
// It will always return an error to detect redirects and to restart the
// latency measurement
func NoRedirectsPolicy(req *http.Request, via []*http.Request) error {
	return LatencyRedirectError(fmt.Sprintf("%v", req.URL))
}

// Performans the LatencyRequest and returns a LatencyResponse
func (l *LatencyRequest) Execute() (resp *LatencyResponse, err error) {

	client := &http.Client{
		Timeout:       l.Timeout,
		CheckRedirect: NoRedirectsPolicy,
	}

	if len(l.Redirects) > 0 {
		// Parse redirect target
		newUrl, err := url.Parse(l.Redirects[len(l.Redirects)-1])
		if err != nil {
			return nil, err
		}
		// Adapt request
		l.Request.URL = newUrl
		l.Request.Host = newUrl.Host
	} else {
		l.Redirects = append(l.Redirects, l.Request.URL.String())
	}
	l.Request.Header.Set("Cache control", "no-cache")

	timeStart := time.Now()
	response, err := client.Do(l.Request)
	if err != nil {
		// Redirect error
		if strings.Contains(err.Error(), ErrLatencyRedirect) {
			if loc := response.Header.Get("Location"); loc != "" {
				l.Redirects = append(l.Redirects, loc)
				return l.Execute()
			} else {
				return nil, errors.New("No redirect location given!")
			}
		} else {
			return nil, err
		}
	}
	latency := time.Since(timeStart)

	response.Body.Close()

	resp = &LatencyResponse{response, latency, l.Redirects}
	return resp, nil
}

// LatencyResponse is a wrapper for http.Response.
// It contains the latency that was measured
type LatencyResponse struct {
	*http.Response
	Latency   time.Duration
	Redirects []string
}

func (l *LatencyResponse) String() string {
	if len(l.Redirects) > 1 {
		return fmt.Sprintf("%v %v %v", l.Request.Method,
			strings.Join(l.Redirects, " -> "), l.Latency)
	}
	return fmt.Sprintf("%v %v %v", l.Request.Method, l.Request.URL, l.Latency)
}
<commit_msg>Godoc fix<commit_after>/*
MAF Latency Package

Measure the latency of HTTP requests.

Usage

First specify a new http.Request like:

	req, _ := http.NewRequest("GET", "http://google.com", nil)

To use this request with the latency library you have to embed it in a
LatencyRequest and specify a timeout interval. Example:

	request := LatencyRequest{
		Request: req,
		Timeout: 10 * time.Second,
	}

To execute the latency measurement, simple call the Execute() method:

	response, err := request.Execute()

Note that the latency measurement follows redirects, but only measures the time
of the last non-redirecting request. The call will return a LatencyResponse
object.
*/
package latency

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"
)

var (
	ErrLatencyRedirect = "Redirect in request"
)

// LatencyRequest is a wrapper for http.Request.
// You can specify an additonal timeout
type LatencyRequest struct {
	*http.Request
	Timeout   time.Duration
	Redirects []string
}

func (l *LatencyRequest) String() string {
	return fmt.Sprintf("%v %v, timout is %v", l.Method, l.URL, l.Timeout)
}

// LatencyRedirectError is an error type used, if redirection occured during
// latency measurement
type LatencyRedirectError string

func (l LatencyRedirectError) Error() string {
	return ErrLatencyRedirect
}

// Redirect policy for http.Client
// It will always return an error to detect redirects and to restart the
// latency measurement
func NoRedirectsPolicy(req *http.Request, via []*http.Request) error {
	return LatencyRedirectError(fmt.Sprintf("%v", req.URL))
}

// Performans the LatencyRequest and returns a LatencyResponse
func (l *LatencyRequest) Execute() (resp *LatencyResponse, err error) {

	client := &http.Client{
		Timeout:       l.Timeout,
		CheckRedirect: NoRedirectsPolicy,
	}

	if len(l.Redirects) > 0 {
		// Parse redirect target
		newUrl, err := url.Parse(l.Redirects[len(l.Redirects)-1])
		if err != nil {
			return nil, err
		}
		// Adapt request
		l.Request.URL = newUrl
		l.Request.Host = newUrl.Host
	} else {
		l.Redirects = append(l.Redirects, l.Request.URL.String())
	}
	l.Request.Header.Set("Cache control", "no-cache")

	timeStart := time.Now()
	response, err := client.Do(l.Request)
	if err != nil {
		// Redirect error
		if strings.Contains(err.Error(), ErrLatencyRedirect) {
			if loc := response.Header.Get("Location"); loc != "" {
				l.Redirects = append(l.Redirects, loc)
				return l.Execute()
			} else {
				return nil, errors.New("No redirect location given!")
			}
		} else {
			return nil, err
		}
	}
	latency := time.Since(timeStart)

	response.Body.Close()

	resp = &LatencyResponse{response, latency, l.Redirects}
	return resp, nil
}

// LatencyResponse is a wrapper for http.Response.
// It contains the latency that was measured
type LatencyResponse struct {
	*http.Response
	Latency   time.Duration
	Redirects []string
}

func (l *LatencyResponse) String() string {
	if len(l.Redirects) > 1 {
		return fmt.Sprintf("%v %v %v", l.Request.Method,
			strings.Join(l.Redirects, " -> "), l.Latency)
	}
	return fmt.Sprintf("%v %v %v", l.Request.Method, l.Request.URL, l.Latency)
}
<|endoftext|>


<commit_before>// A simple WebSocket proxy (WSPR) that takes in a Veyron RPC message, encoded in JSON
// and stored in a WebSocket message, and sends it to the specified Veyron
// endpoint.
//
// Input arguments must be provided as a JSON message in the following format:
//
// {
//   "Address" : String, //EndPoint Address
//   "Name" : String, //Service Name
//   "Method"   : String, //Method Name
//   "InArgs"     : { "ArgName1" : ArgVal1, "ArgName2" : ArgVal2, ... },
//   "IsStreaming" : true/false
// }
//
package wspr

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	_ "net/http/pprof"
	"sync"
	"time"

	"veyron.io/veyron/veyron2"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/rt"
	"veyron.io/veyron/veyron2/security"
	"veyron.io/veyron/veyron2/vlog"

	veyron_identity "veyron.io/veyron/veyron/services/identity"
	"veyron.io/wspr/veyron/services/wsprd/identity"
)

const (
	pingInterval = 50 * time.Second              // how often the server pings the client.
	pongTimeout  = pingInterval + 10*time.Second // maximum wait for pong.
)

type wsprConfig struct {
	MounttableRoot []string
}

type WSPR struct {
	mu             sync.Mutex
	tlsCert        *tls.Certificate
	rt             veyron2.Runtime
	httpPort       int // HTTP port for WSPR to serve on. Port rather than address to discourage serving in a way that isn't local.
	logger         vlog.Logger
	listenSpec     ipc.ListenSpec
	identdEP       string
	idManager      *identity.IDManager
	blesserService veyron_identity.OAuthBlesser
	pipes          map[*http.Request]*pipe
}

var logger vlog.Logger

func readFromRequest(r *http.Request) (*bytes.Buffer, error) {
	var buf bytes.Buffer
	if readBytes, err := io.Copy(&buf, r.Body); err != nil {
		return nil, fmt.Errorf("error copying message out of request: %v", err)
	} else if wantBytes := r.ContentLength; readBytes != wantBytes {
		return nil, fmt.Errorf("read %d bytes, wanted %d", readBytes, wantBytes)
	}
	return &buf, nil
}

func setAccessControl(w http.ResponseWriter) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
}

// Starts the proxy and listens for requests. This method is blocking.
func (ctx WSPR) Run() {
	// Bind to the OAuth Blesser service
	blesserService, err := veyron_identity.BindOAuthBlesser(ctx.identdEP)
	if err != nil {
		log.Fatalf("Failed to bind to identity service at %v: %v", ctx.identdEP, err)
	}
	ctx.blesserService = blesserService

	// HTTP routes
	http.HandleFunc("/debug", ctx.handleDebug)
	http.HandleFunc("/create-account", ctx.handleCreateAccount)
	http.HandleFunc("/assoc-account", ctx.handleAssocAccount)
	http.HandleFunc("/ws", ctx.handleWS)
	// Everything else is a 404.
	// Note: the pattern "/" matches all paths not matched by other
	// registered patterns, not just the URL with Path == "/".'
	// (http://golang.org/pkg/net/http/#ServeMux)
	http.Handle("/", http.NotFoundHandler())
	ctx.logger.VI(1).Infof("Listening at port %d.", ctx.httpPort)
	httpErr := http.ListenAndServe(fmt.Sprintf("127.0.0.1:%d", ctx.httpPort), nil)
	if httpErr != nil {
		log.Fatalf("Failed to HTTP serve: %s", httpErr)
	}
}

func (ctx WSPR) Shutdown() {
	ctx.rt.Cleanup()
}

func (ctx WSPR) CleanUpPipe(req *http.Request) {
	ctx.mu.Lock()
	defer ctx.mu.Unlock()
	delete(ctx.pipes, req)
}

// Creates a new WebSocket Proxy object.
func NewWSPR(httpPort int, listenSpec ipc.ListenSpec, identdEP string, opts ...veyron2.ROpt) *WSPR {
	if listenSpec.Proxy == "" {
		log.Fatalf("a veyron proxy must be set")
	}
	if identdEP == "" {
		log.Fatalf("an identd server must be set")
	}

	newrt, err := rt.New(opts...)
	if err != nil {
		log.Fatalf("rt.New failed: %s", err)
	}

	// TODO(nlacasse, bjornick) use a serializer that can actually persist.
	idManager, err := identity.NewIDManager(newrt, &identity.InMemorySerializer{})
	if err != nil {
		log.Fatalf("identity.NewIDManager failed: %s", err)
	}

	return &WSPR{
		httpPort:   httpPort,
		listenSpec: listenSpec,
		identdEP:   identdEP,
		rt:         newrt,
		logger:     newrt.Logger(),
		idManager:  idManager,
		pipes:      map[*http.Request]*pipe{},
	}
}

// HTTP Handlers

func (ctx WSPR) handleDebug(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		fmt.Fprintf(w, "")
		return
	}
	w.Header().Set("Content-Type", "text/html")
	w.Write([]byte(`

<title>/debug</title>

/debug/pprof

`))
}

func (ctx WSPR) handleWS(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
		return
	}
	ctx.logger.VI(0).Info("Creating a new websocket")
	p := newPipe(w, r, &ctx, nil)

	if p == nil {
		return
	}
	ctx.mu.Lock()
	defer ctx.mu.Unlock()
	ctx.pipes[r] = p
}

// Structs for marshalling input/output to create-account route.
type createAccountInput struct {
	AccessToken string `json:"access_token"`
}

type createAccountOutput struct {
	Names []string `json:"names"`
}

// Handler for creating an account in the identity manager.
// A valid OAuth2 access token must be supplied in the request body. That
// access token is exchanged for a blessing from the identd server.  A new
// privateID is then derived from WSPR's privateID and the blessing. That
// privateID is stored in the identity manager. The name of the new privateID
// is returned to the client.
func (ctx WSPR) handleCreateAccount(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
		return
	}

	// Parse request body.
	var data createAccountInput
	if err := json.NewDecoder(r.Body).Decode(&data); err != nil {
		msg := fmt.Sprintf("Error parsing body: %v", err)
		ctx.logger.Error(msg)
		http.Error(w, msg, http.StatusBadRequest)
	}

	// Get a blessing for the access token from identity server.
	rctx, cancel := ctx.rt.NewContext().WithTimeout(time.Minute)
	defer cancel()
	blessingAny, _, err := ctx.blesserService.BlessUsingAccessToken(rctx, data.AccessToken)
	if err != nil {
		msg := fmt.Sprintf("Error getting blessing for access token: %v", err)
		ctx.logger.Error(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}
	blessing := blessingAny.(security.PublicID)

	// Derive a new identity from the runtime's identity and the blessing.
	identity, err := ctx.rt.Identity().Derive(blessing)
	if err != nil {
		msg := fmt.Sprintf("Error deriving identity: %v", err)
		ctx.logger.Error(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}

	for _, name := range blessing.Names() {
		// Store identity in identity manager.
		if err := ctx.idManager.AddAccount(name, identity); err != nil {
			msg := fmt.Sprintf("Error storing identity: %v", err)
			ctx.logger.Error(msg)
			http.Error(w, msg, http.StatusBadRequest)
			return
		}
	}

	// Return the names to the client.
	out := createAccountOutput{
		Names: blessing.Names(),
	}
	outJson, err := json.Marshal(out)
	if err != nil {
		msg := fmt.Sprintf("Error mashalling names: %v", err)
		ctx.logger.Error(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}

	// Success.
	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, string(outJson))
}

// Struct for marshalling input to assoc-account route.
type assocAccountInput struct {
	Name   string `json:"name"`
	Origin string `json:"origin"`
}

// Handler for associating an existing privateID with an origin.
func (ctx WSPR) handleAssocAccount(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
		return
	}

	// Parse request body.
	var data assocAccountInput
	if err := json.NewDecoder(r.Body).Decode(&data); err != nil {
		http.Error(w, fmt.Sprintf("Error parsing body: %v", err), http.StatusBadRequest)
	}

	// Store the origin.
	// TODO(nlacasse, bjornick): determine what the caveats should be.
	if err := ctx.idManager.AddOrigin(data.Origin, data.Name, nil); err != nil {
		http.Error(w, fmt.Sprintf("Error associating account: %v", err), http.StatusBadRequest)
		return
	}

	// Success.
	fmt.Fprintf(w, "")
}
<commit_msg>veyron.io/wspr: Removing unused function.<commit_after>// A simple WebSocket proxy (WSPR) that takes in a Veyron RPC message, encoded in JSON
// and stored in a WebSocket message, and sends it to the specified Veyron
// endpoint.
//
// Input arguments must be provided as a JSON message in the following format:
//
// {
//   "Address" : String, //EndPoint Address
//   "Name" : String, //Service Name
//   "Method"   : String, //Method Name
//   "InArgs"     : { "ArgName1" : ArgVal1, "ArgName2" : ArgVal2, ... },
//   "IsStreaming" : true/false
// }
//
package wspr

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	_ "net/http/pprof"
	"sync"
	"time"

	"veyron.io/veyron/veyron2"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/rt"
	"veyron.io/veyron/veyron2/security"
	"veyron.io/veyron/veyron2/vlog"

	veyron_identity "veyron.io/veyron/veyron/services/identity"
	"veyron.io/wspr/veyron/services/wsprd/identity"
)

const (
	pingInterval = 50 * time.Second              // how often the server pings the client.
	pongTimeout  = pingInterval + 10*time.Second // maximum wait for pong.
)

type wsprConfig struct {
	MounttableRoot []string
}

type WSPR struct {
	mu             sync.Mutex
	tlsCert        *tls.Certificate
	rt             veyron2.Runtime
	httpPort       int // HTTP port for WSPR to serve on. Port rather than address to discourage serving in a way that isn't local.
	logger         vlog.Logger
	listenSpec     ipc.ListenSpec
	identdEP       string
	idManager      *identity.IDManager
	blesserService veyron_identity.OAuthBlesser
	pipes          map[*http.Request]*pipe
}

var logger vlog.Logger

func readFromRequest(r *http.Request) (*bytes.Buffer, error) {
	var buf bytes.Buffer
	if readBytes, err := io.Copy(&buf, r.Body); err != nil {
		return nil, fmt.Errorf("error copying message out of request: %v", err)
	} else if wantBytes := r.ContentLength; readBytes != wantBytes {
		return nil, fmt.Errorf("read %d bytes, wanted %d", readBytes, wantBytes)
	}
	return &buf, nil
}

// Starts the proxy and listens for requests. This method is blocking.
func (ctx WSPR) Run() {
	// Bind to the OAuth Blesser service
	blesserService, err := veyron_identity.BindOAuthBlesser(ctx.identdEP)
	if err != nil {
		log.Fatalf("Failed to bind to identity service at %v: %v", ctx.identdEP, err)
	}
	ctx.blesserService = blesserService

	// HTTP routes
	http.HandleFunc("/debug", ctx.handleDebug)
	http.HandleFunc("/create-account", ctx.handleCreateAccount)
	http.HandleFunc("/assoc-account", ctx.handleAssocAccount)
	http.HandleFunc("/ws", ctx.handleWS)
	// Everything else is a 404.
	// Note: the pattern "/" matches all paths not matched by other
	// registered patterns, not just the URL with Path == "/".'
	// (http://golang.org/pkg/net/http/#ServeMux)
	http.Handle("/", http.NotFoundHandler())
	ctx.logger.VI(1).Infof("Listening at port %d.", ctx.httpPort)
	httpErr := http.ListenAndServe(fmt.Sprintf("127.0.0.1:%d", ctx.httpPort), nil)
	if httpErr != nil {
		log.Fatalf("Failed to HTTP serve: %s", httpErr)
	}
}

func (ctx WSPR) Shutdown() {
	ctx.rt.Cleanup()
}

func (ctx WSPR) CleanUpPipe(req *http.Request) {
	ctx.mu.Lock()
	defer ctx.mu.Unlock()
	delete(ctx.pipes, req)
}

// Creates a new WebSocket Proxy object.
func NewWSPR(httpPort int, listenSpec ipc.ListenSpec, identdEP string, opts ...veyron2.ROpt) *WSPR {
	if listenSpec.Proxy == "" {
		log.Fatalf("a veyron proxy must be set")
	}
	if identdEP == "" {
		log.Fatalf("an identd server must be set")
	}

	newrt, err := rt.New(opts...)
	if err != nil {
		log.Fatalf("rt.New failed: %s", err)
	}

	// TODO(nlacasse, bjornick) use a serializer that can actually persist.
	idManager, err := identity.NewIDManager(newrt, &identity.InMemorySerializer{})
	if err != nil {
		log.Fatalf("identity.NewIDManager failed: %s", err)
	}

	return &WSPR{
		httpPort:   httpPort,
		listenSpec: listenSpec,
		identdEP:   identdEP,
		rt:         newrt,
		logger:     newrt.Logger(),
		idManager:  idManager,
		pipes:      map[*http.Request]*pipe{},
	}
}

// HTTP Handlers

func (ctx WSPR) handleDebug(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		fmt.Fprintf(w, "")
		return
	}
	w.Header().Set("Content-Type", "text/html")
	w.Write([]byte(`

<title>/debug</title>

/debug/pprof

`))
}

func (ctx WSPR) handleWS(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
		return
	}
	ctx.logger.VI(0).Info("Creating a new websocket")
	p := newPipe(w, r, &ctx, nil)

	if p == nil {
		return
	}
	ctx.mu.Lock()
	defer ctx.mu.Unlock()
	ctx.pipes[r] = p
}

// Structs for marshalling input/output to create-account route.
type createAccountInput struct {
	AccessToken string `json:"access_token"`
}

type createAccountOutput struct {
	Names []string `json:"names"`
}

// Handler for creating an account in the identity manager.
// A valid OAuth2 access token must be supplied in the request body. That
// access token is exchanged for a blessing from the identd server.  A new
// privateID is then derived from WSPR's privateID and the blessing. That
// privateID is stored in the identity manager. The name of the new privateID
// is returned to the client.
func (ctx WSPR) handleCreateAccount(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
		return
	}

	// Parse request body.
	var data createAccountInput
	if err := json.NewDecoder(r.Body).Decode(&data); err != nil {
		msg := fmt.Sprintf("Error parsing body: %v", err)
		ctx.logger.Error(msg)
		http.Error(w, msg, http.StatusBadRequest)
	}

	// Get a blessing for the access token from identity server.
	rctx, cancel := ctx.rt.NewContext().WithTimeout(time.Minute)
	defer cancel()
	blessingAny, _, err := ctx.blesserService.BlessUsingAccessToken(rctx, data.AccessToken)
	if err != nil {
		msg := fmt.Sprintf("Error getting blessing for access token: %v", err)
		ctx.logger.Error(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}
	blessing := blessingAny.(security.PublicID)

	// Derive a new identity from the runtime's identity and the blessing.
	identity, err := ctx.rt.Identity().Derive(blessing)
	if err != nil {
		msg := fmt.Sprintf("Error deriving identity: %v", err)
		ctx.logger.Error(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}

	for _, name := range blessing.Names() {
		// Store identity in identity manager.
		if err := ctx.idManager.AddAccount(name, identity); err != nil {
			msg := fmt.Sprintf("Error storing identity: %v", err)
			ctx.logger.Error(msg)
			http.Error(w, msg, http.StatusBadRequest)
			return
		}
	}

	// Return the names to the client.
	out := createAccountOutput{
		Names: blessing.Names(),
	}
	outJson, err := json.Marshal(out)
	if err != nil {
		msg := fmt.Sprintf("Error mashalling names: %v", err)
		ctx.logger.Error(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}

	// Success.
	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, string(outJson))
}

// Struct for marshalling input to assoc-account route.
type assocAccountInput struct {
	Name   string `json:"name"`
	Origin string `json:"origin"`
}

// Handler for associating an existing privateID with an origin.
func (ctx WSPR) handleAssocAccount(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed.", http.StatusMethodNotAllowed)
		return
	}

	// Parse request body.
	var data assocAccountInput
	if err := json.NewDecoder(r.Body).Decode(&data); err != nil {
		http.Error(w, fmt.Sprintf("Error parsing body: %v", err), http.StatusBadRequest)
	}

	// Store the origin.
	// TODO(nlacasse, bjornick): determine what the caveats should be.
	if err := ctx.idManager.AddOrigin(data.Origin, data.Name, nil); err != nil {
		http.Error(w, fmt.Sprintf("Error associating account: %v", err), http.StatusBadRequest)
		return
	}

	// Success.
	fmt.Fprintf(w, "")
}
<|endoftext|>


<commit_before>package simplehstore

import (
	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/lib/pq"
)

// NewKeyValue creates a new KeyValue struct, for storing key/value pairs.
func NewKeyValue(host *Host, name string) (*KeyValue, error) {
	kv := &KeyValue{host, name}

	// Create extension hstore
	query := "CREATE EXTENSION hstore"
	if Verbose {
		fmt.Println(query)
	}
	// Ignore erors if this is already created
	kv.host.db.Exec(query)

	query = fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (attr hstore)", pq.QuoteIdentifier(kvPrefix+kv.table))
	if _, err := kv.host.db.Exec(query); err != nil {
		return nil, err
	}
	if Verbose {
		log.Println("Created HSTORE table " + pq.QuoteIdentifier(kvPrefix+kv.table) + " in database " + host.dbname)
	}
	return kv, nil
}

// CreateIndexTable creates an INDEX table for this key/value, that may speed up lookups
func (kv *KeyValue) CreateIndexTable() error {
	// strip double quotes from kv.table and add _idx at the end
	indexTableName := strings.TrimSuffix(strings.TrimPrefix(kv.table, "\""), "\"") + "_idx"
	query := fmt.Sprintf("CREATE INDEX %q ON %s USING GIN (attr)", indexTableName, kv.table)
	if Verbose {
		fmt.Println(query)
	}
	_, err := kv.host.db.Exec(query)
	return err
}

// RemoveIndexTable removes the INDEX table for this key/value
func (kv *KeyValue) RemoveIndexTable() error {
	// strip double quotes from kv.table and add _idx at the end
	indexTableName := strings.TrimSuffix(strings.TrimPrefix(kv.table, "\""), "\"") + "_idx"
	query := fmt.Sprintf("DROP INDEX %q", indexTableName)
	if Verbose {
		fmt.Println(query)
	}
	_, err := kv.host.db.Exec(query)
	return err
}

// Set a key and value
func (kv *KeyValue) Set(key, value string) error {
	if !kv.host.rawUTF8 {
		Encode(&value)
	}
	if _, err := kv.Get(key); err != nil {
		// Key does not exist, create it
		_, err = kv.host.db.Exec(fmt.Sprintf("INSERT INTO %s (attr) VALUES ('\"%s\"=>\"%s\"')", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key), escapeSingleQuotes(value)))
		return err
	}
	// Key exists, update the value
	_, err := kv.host.db.Exec(fmt.Sprintf("UPDATE %s SET attr = attr || '\"%s\"=>\"%s\"' :: hstore", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key), escapeSingleQuotes(value)))
	return err
}

func (kv *KeyValue) insert(key, encodedValue string) (int64, error) {
	query := fmt.Sprintf("INSERT INTO %s (attr) VALUES ('\"%s\"=>\"%s\"')", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key), escapeSingleQuotes(encodedValue))
	if Verbose {
		fmt.Println(query)
	}
	result, err := kv.host.db.Exec(query)
	if Verbose {
		log.Println("Inserted row into: "+kv.table+" err? ", err)
	}
	n, _ := result.RowsAffected()
	return n, err
}

func (kv *KeyValue) update(key, encodedValue string) (int64, error) {
	query := fmt.Sprintf("UPDATE %s SET attr = attr || '\"%s\"=>\"%s\"' :: hstore", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key), escapeSingleQuotes(encodedValue))
	if Verbose {
		fmt.Println(query)
	}
	result, err := kv.host.db.Exec(query)
	if Verbose {
		log.Println("Updated row in: "+kv.table+" err? ", err)
	}
	n, _ := result.RowsAffected()
	return n, err
}

// SetCheck will set a value in a hashmap given the element id (for instance a user id) and the key (for instance "password")
// Returns true if the key already existed.
func (kv *KeyValue) SetCheck(key, value string) (bool, error) {
	if !kv.host.rawUTF8 {
		Encode(&value)
	}
	encodedValue := value
	// First try updating the key/values
	n, err := kv.update(key, encodedValue)
	if err != nil {
		return false, err
	}
	// If no rows are affected (SELECTED) by the update, try inserting a row instead
	if n == 0 {
		n, err = kv.insert(key, encodedValue)
		if err != nil {
			return false, err
		}
		if n == 0 {
			return false, errors.New("could not update or insert any rows")
		}
		return false, nil
	}
	// success, and the key already existed
	return true, nil
}

// Get a value given a key
func (kv *KeyValue) Get(key string) (string, error) {
	rows, err := kv.host.db.Query(fmt.Sprintf("SELECT attr -> '%s' FROM %s", escapeSingleQuotes(key), pq.QuoteIdentifier(kvPrefix+kv.table)))
	if err != nil {
		return "", err
	}
	if rows == nil {
		return "", errors.New("KeyValue Get returned no rows for key " + key)
	}
	defer rows.Close()
	var value string
	// Get the value. Should only loop once.
	counter := 0
	for rows.Next() {
		err = rows.Scan(&value)
		if err != nil {
			return "", err
		}
		counter++
	}
	if err := rows.Err(); err != nil {
		return "", err
	}
	if counter != 1 {
		return "", errors.New("Wrong number of keys in KeyValue table: " + kvPrefix + kv.table)
	}
	if !kv.host.rawUTF8 {
		Decode(&value)
	}
	return value, nil
}

// Inc increases the value of a key and returns the new value.
// Returns "1" if no previous value is found.
func (kv *KeyValue) Inc(key string) (string, error) {
	// Retrieve the current value, if any
	num := 0
	// See if we can fetch an existing value.
	if val, err := kv.Get(key); err == nil { // success
		// See if we can convert the value to a number.
		if converted, errConv := strconv.Atoi(val); errConv == nil { // success
			num = converted
		}
	} else {
		// The key does not exist, create a new one.
		// This is to reflect the behavior of INCR in Redis.
		NewKeyValue(kv.host, kv.table)
	}
	// Num is now either 0 or the previous numeric value
	num++
	// Convert the new value to a string
	val := strconv.Itoa(num)
	// Store the new number
	if err := kv.Set(key, val); err != nil {
		// Saving the value failed
		return "0", err
	}
	// Success
	return val, nil
}

// Dec increases the value of a key and returns the new value.
// Returns "1" if no previous value is found.
func (kv *KeyValue) Dec(key string) (string, error) {
	// Retrieve the current value, if any
	num := 0
	// See if we can fetch an existing value. NOTE: "== nil"
	if val, err := kv.Get(key); err == nil {
		// See if we can convert the value to a number. NOTE: "== nil"
		if converted, errConv := strconv.Atoi(val); errConv == nil {
			num = converted
		}
	} else {
		// The key does not exist, create a new one.
		NewKeyValue(kv.host, kv.table)
	}
	// Num is now either 0 or the previous numeric value
	num--
	// Convert the new value to a string
	val := strconv.Itoa(num)
	// Store the new number
	if err := kv.Set(key, val); err != nil {
		// Saving the value failed
		return "0", err
	}
	// Success
	return val, nil
}

// Del removes the given key
func (kv *KeyValue) Del(key string) error {
	_, err := kv.host.db.Exec(fmt.Sprintf("UPDATE %s SET attr = delete(attr, '%s')", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key)))
	return err
}

// Remove this key/value
func (kv *KeyValue) Remove() error {
	// Remove the table
	_, err := kv.host.db.Exec(fmt.Sprintf("DROP TABLE %s", pq.QuoteIdentifier(kvPrefix+kv.table)))
	return err
}

// Clear this key/value
func (kv *KeyValue) Clear() error {
	// Truncate the table
	_, err := kv.host.db.Exec(fmt.Sprintf("TRUNCATE TABLE %s", pq.QuoteIdentifier(kvPrefix+kv.table)))
	return err
}
<commit_msg>Use sql.NullString + add transactions<commit_after>package simplehstore

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/lib/pq"
)

// KeyValue is a hash map with a key and a value, stored in PostgreSQL
type KeyValue dbDatastructure

// NewKeyValue creates a new KeyValue struct, for storing key/value pairs.
func NewKeyValue(host *Host, name string) (*KeyValue, error) {
	kv := &KeyValue{host, name}

	// Create extension hstore
	query := "CREATE EXTENSION hstore"
	// Ignore erors if this is already created
	kv.host.db.Exec(query)

	query = fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (attr hstore)", pq.QuoteIdentifier(kvPrefix+kv.table))
	if _, err := kv.host.db.Exec(query); err != nil {
		return nil, err
	}
	if Verbose {
		log.Println("Created HSTORE table " + pq.QuoteIdentifier(kvPrefix+kv.table) + " in database " + host.dbname)
	}
	return kv, nil
}

// CreateIndexTable creates an INDEX table for this key/value, that may speed up lookups
func (kv *KeyValue) CreateIndexTable() error {
	// strip double quotes from kv.table and add _idx at the end
	indexTableName := strings.TrimSuffix(strings.TrimPrefix(kv.table, "\""), "\"") + "_idx"
	query := fmt.Sprintf("CREATE INDEX %q ON %s USING GIN (attr)", indexTableName, kv.table)
	if Verbose {
		fmt.Println(query)
	}
	_, err := kv.host.db.Exec(query)
	return err
}

// RemoveIndexTable removes the INDEX table for this key/value
func (kv *KeyValue) RemoveIndexTable() error {
	// strip double quotes from kv.table and add _idx at the end
	indexTableName := strings.TrimSuffix(strings.TrimPrefix(kv.table, "\""), "\"") + "_idx"
	query := fmt.Sprintf("DROP INDEX %q", indexTableName)
	if Verbose {
		fmt.Println(query)
	}
	_, err := kv.host.db.Exec(query)
	return err
}

// insert a new key+value in the current KeyValue table
func (kv *KeyValue) insert(key, encodedValue string) (int64, error) {
	// Try inserting
	query := fmt.Sprintf("INSERT INTO %s (attr) VALUES ('\"%s\"=>\"%s\"')", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key), escapeSingleQuotes(encodedValue))
	if Verbose {
		fmt.Println(query)
	}
	result, err := kv.host.db.Exec(query)
	if Verbose {
		log.Println("Inserted row into: "+kv.table+" err? ", err)
	}
	n, _ := result.RowsAffected()
	return n, err
}

// insert a new key+value in the current KeyValue table, as part of a transaction
func (kv *KeyValue) insertWithTransaction(ctx context.Context, transaction *sql.Tx, key, encodedValue string) (int64, error) {
	// Try inserting
	query := fmt.Sprintf("INSERT INTO %s (attr) VALUES ('\"%s\"=>\"%s\"')", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key), escapeSingleQuotes(encodedValue))
	if Verbose {
		fmt.Println(query)
	}
	result, err := transaction.ExecContext(ctx, query)
	if Verbose {
		log.Println("Inserted row into: "+kv.table+" err? ", err)
	}
	n, _ := result.RowsAffected()
	return n, err
}

// update a value in the current KeyValue table
func (kv *KeyValue) update(key, encodedValue string) (int64, error) {
	// Try updating
	query := fmt.Sprintf("UPDATE %s SET attr = attr || '\"%s\"=>\"%s\"' :: hstore", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key), escapeSingleQuotes(encodedValue))
	if Verbose {
		fmt.Println(query)
	}
	result, err := kv.host.db.Exec(query)
	if Verbose {
		log.Println("Updated row in: "+kv.table+" err? ", err)
	}
	if result == nil {
		return 0, fmt.Errorf("no result when trying to update %s with a value", key)
	}
	n, _ := result.RowsAffected()
	return n, err
}

// update a value in the current KeyValue table, as part of a transaction
func (kv *KeyValue) updateWithTransaction(ctx context.Context, transaction *sql.Tx, key, encodedValue string) (int64, error) {
	// Try updating
	query := fmt.Sprintf("UPDATE %s SET attr = attr || '\"%s\"=>\"%s\"' :: hstore", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key), escapeSingleQuotes(encodedValue))
	if Verbose {
		fmt.Println(query)
	}
	result, err := transaction.ExecContext(ctx, query)
	if Verbose {
		log.Println("Updated row in: "+kv.table+" err? ", err)
	}
	if result == nil {
		return 0, fmt.Errorf("no result when trying to update %s with a value", key)
	}
	n, _ := result.RowsAffected()
	return n, err
}

// Set a key and value
func (kv *KeyValue) Set(key, value string) error {
	if !kv.host.rawUTF8 {
		Encode(&value)
	}
	encodedValue := value
	// First try updating the key/values
	n, err := kv.update(key, encodedValue)
	if err != nil {
		return err
	}
	// If no rows are affected (SELECTED) by the update, try inserting a row instead
	if n == 0 {
		n, err = kv.insert(key, encodedValue)
		if err != nil {
			return err
		}
		if n == 0 {
			return errors.New("could not update or insert any rows")
		}
	}
	// success
	return nil
}

// set a key and value, as part of a transaction
func (kv *KeyValue) setWithTransaction(ctx context.Context, transaction *sql.Tx, key, value string) error {
	if !kv.host.rawUTF8 {
		Encode(&value)
	}
	encodedValue := value
	// First try updating the key/values
	n, err := kv.updateWithTransaction(ctx, transaction, key, encodedValue)
	if err != nil {
		return err
	}
	// If no rows are affected (SELECTED) by the update, try inserting a row instead
	if n == 0 {
		n, err = kv.insertWithTransaction(ctx, transaction, key, encodedValue)
		if err != nil {
			return err
		}
		if n == 0 {
			return errors.New("could not update or insert any rows")
		}
	}
	// success
	return nil
}

// SetCheck will set a value in a hashmap given the element id (for instance a user id) and the key (for instance "password")
// Returns true if the key already existed.
func (kv *KeyValue) SetCheck(key, value string) (bool, error) {
	if !kv.host.rawUTF8 {
		Encode(&value)
	}
	encodedValue := value
	// First try updating the key/values
	n, err := kv.update(key, encodedValue)
	if err != nil {
		return false, err
	}
	// If no rows are affected (SELECTED) by the update, try inserting a row instead
	if n == 0 {
		n, err = kv.insert(key, encodedValue)
		if err != nil {
			return false, err
		}
		if n == 0 {
			return false, errors.New("could not update or insert any rows")
		}
		return false, nil
	}
	// success, and the key already existed
	return true, nil
}

// Get a value given a key
func (kv *KeyValue) Get(key string) (string, error) {
	rows, err := kv.host.db.Query(fmt.Sprintf("SELECT attr -> '%s' FROM %s", escapeSingleQuotes(key), pq.QuoteIdentifier(kvPrefix+kv.table)))
	if err != nil {
		return "", fmt.Errorf("KeyValue.Get: query error: %s", err)
	}
	if rows == nil {
		return "", errors.New("KeyValue.Get: no rows for key " + key)
	}
	defer rows.Close()
	var value sql.NullString
	// Get the value. Should only loop once.
	counter := 0
	for rows.Next() {
		err = rows.Scan(&value)
		if err != nil {
			return "", err
		}
		counter++
	}
	if err := rows.Err(); err != nil {
		return "", fmt.Errorf("KeyValue.Get: rows.Err(): %s", err)
	}
	if counter != 1 {
		return "", fmt.Errorf("Wrong number of keys in KeyValue table: %s", kvPrefix+kv.table)
	}
	if counter == 0 {
		return "", fmt.Errorf("No rows")
	}
	s := value.String
	if !kv.host.rawUTF8 {
		Decode(&s)
	}
	return s, nil
}

// Inc increases the value of a key and returns the new value.
// Returns "1" if no previous value is found.
func (kv *KeyValue) Inc(key string) (string, error) {
	// Retrieve the current value, if any
	num := 0
	// See if we can fetch an existing value.
	if val, err := kv.Get(key); err == nil { // success
		// See if we can convert the value to a number.
		if converted, errConv := strconv.Atoi(val); errConv == nil { // success
			num = converted
		}
	} else {
		// The key does not exist, create a new one.
		// This is to reflect the behavior of INCR in Redis.
		NewKeyValue(kv.host, kv.table)
	}
	// Num is now either 0 or the previous numeric value
	num++
	// Convert the new value to a string
	val := strconv.Itoa(num)
	// Store the new number
	if err := kv.Set(key, val); err != nil {
		// Saving the value failed
		return "0", err
	}
	// Success
	return val, nil
}

// Dec increases the value of a key and returns the new value.
// Returns "1" if no previous value is found.
func (kv *KeyValue) Dec(key string) (string, error) {
	// Retrieve the current value, if any
	num := 0
	// See if we can fetch an existing value. NOTE: "== nil"
	if val, err := kv.Get(key); err == nil {
		// See if we can convert the value to a number. NOTE: "== nil"
		if converted, errConv := strconv.Atoi(val); errConv == nil {
			num = converted
		}
	} else {
		// The key does not exist, create a new one.
		NewKeyValue(kv.host, kv.table)
	}
	// Num is now either 0 or the previous numeric value
	num--
	// Convert the new value to a string
	val := strconv.Itoa(num)
	// Store the new number
	if err := kv.Set(key, val); err != nil {
		// Saving the value failed
		return "0", err
	}
	// Success
	return val, nil
}

// Del removes the given key
func (kv *KeyValue) Del(key string) error {
	_, err := kv.host.db.Exec(fmt.Sprintf("UPDATE %s SET attr = delete(attr, '%s')", pq.QuoteIdentifier(kvPrefix+kv.table), escapeSingleQuotes(key)))
	return err
}

// Remove this key/value
func (kv *KeyValue) Remove() error {
	// Remove the table
	_, err := kv.host.db.Exec(fmt.Sprintf("DROP TABLE %s", pq.QuoteIdentifier(kvPrefix+kv.table)))
	return err
}

// Clear this key/value
func (kv *KeyValue) Clear() error {
	// Truncate the table
	_, err := kv.host.db.Exec(fmt.Sprintf("TRUNCATE TABLE %s", pq.QuoteIdentifier(kvPrefix+kv.table)))
	return err
}
<|endoftext|>


<commit_before>package main

import "github.com/akalin/ilium/ilium"
import "encoding/json"
import "flag"
import "fmt"
import "io/ioutil"
import "math/rand"
import "time"
import "os"
import "runtime"
import "runtime/pprof"

func main() {
	numRenderJobs := flag.Int(
		"j", runtime.NumCPU(), "how many render jobs to spawn")

	profilePath := flag.String(
		"p", "", "if non-empty, path to write the cpu profile to")

	outputDir := flag.String(
		"d", "", "if non-empty, directory to prepend to relative "+
			"output paths")

	outputExt := flag.String(
		"x", "", "if non-empty, the extension to append to "+
			"output paths (but before the real extension)")

	flag.Parse()

	if len(*profilePath) > 0 {
		f, err := os.Create(*profilePath)
		if err != nil {
			panic(err)
		}
		defer func() {
			if err := f.Close(); err != nil {
				panic(err)
			}
		}()

		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	runtime.GOMAXPROCS(*numRenderJobs)

	seed := time.Now().UTC().UnixNano()
	rng := rand.New(rand.NewSource(seed))

	if flag.NArg() < 1 {
		fmt.Fprintf(
			os.Stderr, "%s [options] [scene.json...]\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(-1)
	}

	for i := 0; i < flag.NArg(); i++ {
		inputPath := flag.Arg(i)
		fmt.Printf(
			"Processing %s (%d/%d)...\n",
			inputPath, i+1, flag.NArg())
		configBytes, err := ioutil.ReadFile(inputPath)
		if err != nil {
			panic(err)
		}
		var config map[string]interface{}
		err = json.Unmarshal(configBytes, &config)
		if err != nil {
			panic(err)
		}
		sceneConfig := config["scene"].(map[string]interface{})
		scene := ilium.MakeScene(sceneConfig)
		rendererConfig := config["renderer"].(map[string]interface{})
		renderer := ilium.MakeRenderer(rendererConfig)
		renderer.Render(
			*numRenderJobs, rng, &scene, *outputDir, *outputExt)
	}
}
<commit_msg>Add a command-line switch for the RNG seed<commit_after>package main

import "github.com/akalin/ilium/ilium"
import "encoding/json"
import "flag"
import "fmt"
import "io/ioutil"
import "math/rand"
import "time"
import "os"
import "runtime"
import "runtime/pprof"

func main() {
	numRenderJobs := flag.Int(
		"j", runtime.NumCPU(), "how many render jobs to spawn")

	profilePath := flag.String(
		"p", "", "if non-empty, path to write the cpu profile to")

	outputDir := flag.String(
		"d", "", "if non-empty, directory to prepend to relative "+
			"output paths")

	outputExt := flag.String(
		"x", "", "if non-empty, the extension to append to "+
			"output paths (but before the real extension)")

	// This flag, combined with -j=1, can be used to get
	// repeatable renders.
	seed := flag.Int64(
		"s", time.Now().UTC().UnixNano(),
		"the seed to use for the random number generator")

	flag.Parse()

	if len(*profilePath) > 0 {
		f, err := os.Create(*profilePath)
		if err != nil {
			panic(err)
		}
		defer func() {
			if err := f.Close(); err != nil {
				panic(err)
			}
		}()

		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	runtime.GOMAXPROCS(*numRenderJobs)

	rng := rand.New(rand.NewSource(*seed))

	if flag.NArg() < 1 {
		fmt.Fprintf(
			os.Stderr, "%s [options] [scene.json...]\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(-1)
	}

	for i := 0; i < flag.NArg(); i++ {
		inputPath := flag.Arg(i)
		fmt.Printf(
			"Processing %s (%d/%d)...\n",
			inputPath, i+1, flag.NArg())
		configBytes, err := ioutil.ReadFile(inputPath)
		if err != nil {
			panic(err)
		}
		var config map[string]interface{}
		err = json.Unmarshal(configBytes, &config)
		if err != nil {
			panic(err)
		}
		sceneConfig := config["scene"].(map[string]interface{})
		scene := ilium.MakeScene(sceneConfig)
		rendererConfig := config["renderer"].(map[string]interface{})
		renderer := ilium.MakeRenderer(rendererConfig)
		renderer.Render(
			*numRenderJobs, rng, &scene, *outputDir, *outputExt)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package codegen

import (
	"path/filepath"
	"strings"

	"encoding/json"

	"github.com/pkg/errors"
	"github.com/uber/zanzibar/module"
)

// NewDefaultModuleSystem creates a fresh instance of the default zanzibar
// module system (clients, endpoints)
func NewDefaultModuleSystem(h *PackageHelper) (*module.System, error) {
	system := module.NewSystem()
	tmpl, err := NewTemplate(templateDir)

	if err != nil {
		return nil, err
	}

	// Register client module class and type generators
	if err := system.RegisterClass("client", module.Class{
		Directory:         "clients",
		ClassType:         module.MultiModule,
		ClassDependencies: []string{},
	}); err != nil {
		return nil, errors.Wrapf(err, "Error registering client class")
	}

	if err := system.RegisterClassType("client", "http", &HTTPClientGenerator{
		templates:     tmpl,
		packageHelper: h,
	}); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error registering http client class type",
		)
	}

	if err := system.RegisterClassType("client", "tchannel", &TCahnnelClientGenerator{
		templates:     tmpl,
		packageHelper: h,
	}); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error registering TChannel client class type",
		)
	}

	// TODO: Register endpoint module class and type generators
	return system, nil
}

/*
 * HTTP Client Generator
 */

// HTTPClientGenerator generates an instance of a zanzibar http client
type HTTPClientGenerator struct {
	templates     *Template
	packageHelper *PackageHelper
}

// Generate returns the HTTP client generated files as a map of relative file
// path (relative to the target buid directory) to file bytes.
func (generator *HTTPClientGenerator) Generate(
	instance *module.Instance,
) (map[string][]byte, error) {
	// Parse the client config from the endpoint JSON file
	clientConfig, err := readClientConfig(instance.JSONFileRaw)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading HTTP client %s JSON config",
			instance.InstanceName,
		)
	}

	clientSpec, err := NewHTTPClientSpec(
		filepath.Join(
			instance.BaseDirectory,
			instance.Directory,
			instance.JSONFileName,
		),
		clientConfig,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error initializing HTTPClientSpec for %s",
			instance.InstanceName,
		)
	}

	clientMeta := &ClientMeta{
		PackageName:      clientSpec.ModuleSpec.PackageName,
		Services:         clientSpec.ModuleSpec.Services,
		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
		ClientID:         clientSpec.ClientID,
	}

	client, err := generator.templates.execTemplate(
		"http_client.tmpl",
		clientMeta,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing HTTP client template for %s",
			instance.InstanceName,
		)
	}

	structs, err := generator.templates.execTemplate("structs.tmpl", clientMeta)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing HTTP client structs template for %s",
			instance.InstanceName,
		)
	}

	clientDirectory := filepath.Join(
		generator.packageHelper.CodeGenTargetPath(),
		instance.Directory,
	)

	clientFilePath, err := filepath.Rel(clientDirectory, clientSpec.GoFileName)
	if err != nil {
		clientFilePath = clientSpec.GoFileName
	}

	structFilePath, err := filepath.Rel(
		clientDirectory,
		clientSpec.GoStructsFileName,
	)
	if err != nil {
		structFilePath = clientSpec.GoStructsFileName
	}

	// Return the client files
	files := map[string][]byte{}
	files[clientFilePath] = client
	files[structFilePath] = structs
	return files, nil
}

/*
 * TChannel Client Generator
 */

// TCahnnelClientGenerator generates an instance of a zanzibar TChannel client
type TCahnnelClientGenerator struct {
	templates     *Template
	packageHelper *PackageHelper
}

// Generate returns the TChannel client generated files as a map of relative file
// path (relative to the target build directory) to file bytes.
func (generator *TCahnnelClientGenerator) Generate(
	instance *module.Instance,
) (map[string][]byte, error) {
	// Parse the client config from the endpoint JSON file
	clientConfig, err := readClientConfig(instance.JSONFileRaw)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading TChannel client %s JSON config",
			instance.InstanceName,
		)
	}

	clientSpec, err := NewTChannelClientSpec(
		filepath.Join(
			instance.BaseDirectory,
			instance.Directory,
			instance.JSONFileName,
		),
		clientConfig,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error initializing TChannelClientSpec for %s",
			instance.InstanceName,
		)
	}

	clientMeta := &ClientMeta{
		PackageName:      clientSpec.ModuleSpec.PackageName,
		Services:         clientSpec.ModuleSpec.Services,
		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
		ClientID:         clientSpec.ClientID,
	}

	client, err := generator.templates.execTemplate(
		"tchannel_client.tmpl",
		clientMeta,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel client template for %s",
			instance.InstanceName,
		)
	}

	server, err := generator.templates.execTemplate(
		"tchannel_server.tmpl",
		clientMeta,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel server template for %s",
			instance.InstanceName,
		)
	}

	handler, err := generator.templates.execTemplate(
		"tchannel_handler.tmpl",
		clientMeta,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel handler template for %s",
			instance.InstanceName,
		)
	}

	clientDirectory := filepath.Join(
		generator.packageHelper.CodeGenTargetPath(),
		instance.Directory,
	)

	clientFilePath, err := filepath.Rel(clientDirectory, clientSpec.GoFileName)
	if err != nil {
		clientFilePath = clientSpec.GoFileName
	}

	// TODO:(lu) the locations of tchannel server and handler files are tentative
	serverFilePath := strings.TrimRight(clientFilePath, ".go") + "_server.go"
	handlerFilePath := strings.TrimRight(clientFilePath, ".go") + "_handler.go"

	// Return the client files
	files := map[string][]byte{}
	files[clientFilePath] = client
	files[serverFilePath] = server
	files[handlerFilePath] = handler
	return files, nil
}

func readClientConfig(rawConfig []byte) (*clientClassConfig, error) {
	var clientConfig clientClassConfig
	if err := json.Unmarshal(rawConfig, &clientConfig); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading config for HTTP client instance",
		)
	}
	clientConfig.Config["clientId"] = clientConfig.Name
	clientConfig.Config["clientType"] = clientConfig.Type
	return &clientConfig, nil
}
<commit_msg>Return map literal<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package codegen

import (
	"path/filepath"
	"strings"

	"encoding/json"

	"github.com/pkg/errors"
	"github.com/uber/zanzibar/module"
)

// NewDefaultModuleSystem creates a fresh instance of the default zanzibar
// module system (clients, endpoints)
func NewDefaultModuleSystem(h *PackageHelper) (*module.System, error) {
	system := module.NewSystem()
	tmpl, err := NewTemplate(templateDir)

	if err != nil {
		return nil, err
	}

	// Register client module class and type generators
	if err := system.RegisterClass("client", module.Class{
		Directory:         "clients",
		ClassType:         module.MultiModule,
		ClassDependencies: []string{},
	}); err != nil {
		return nil, errors.Wrapf(err, "Error registering client class")
	}

	if err := system.RegisterClassType("client", "http", &HTTPClientGenerator{
		templates:     tmpl,
		packageHelper: h,
	}); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error registering http client class type",
		)
	}

	if err := system.RegisterClassType("client", "tchannel", &TCahnnelClientGenerator{
		templates:     tmpl,
		packageHelper: h,
	}); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error registering TChannel client class type",
		)
	}

	// TODO: Register endpoint module class and type generators
	return system, nil
}

/*
 * HTTP Client Generator
 */

// HTTPClientGenerator generates an instance of a zanzibar http client
type HTTPClientGenerator struct {
	templates     *Template
	packageHelper *PackageHelper
}

// Generate returns the HTTP client generated files as a map of relative file
// path (relative to the target buid directory) to file bytes.
func (generator *HTTPClientGenerator) Generate(
	instance *module.Instance,
) (map[string][]byte, error) {
	// Parse the client config from the endpoint JSON file
	clientConfig, err := readClientConfig(instance.JSONFileRaw)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading HTTP client %s JSON config",
			instance.InstanceName,
		)
	}

	clientSpec, err := NewHTTPClientSpec(
		filepath.Join(
			instance.BaseDirectory,
			instance.Directory,
			instance.JSONFileName,
		),
		clientConfig,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error initializing HTTPClientSpec for %s",
			instance.InstanceName,
		)
	}

	clientMeta := &ClientMeta{
		PackageName:      clientSpec.ModuleSpec.PackageName,
		Services:         clientSpec.ModuleSpec.Services,
		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
		ClientID:         clientSpec.ClientID,
	}

	client, err := generator.templates.execTemplate(
		"http_client.tmpl",
		clientMeta,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing HTTP client template for %s",
			instance.InstanceName,
		)
	}

	structs, err := generator.templates.execTemplate("structs.tmpl", clientMeta)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing HTTP client structs template for %s",
			instance.InstanceName,
		)
	}

	clientDirectory := filepath.Join(
		generator.packageHelper.CodeGenTargetPath(),
		instance.Directory,
	)

	clientFilePath, err := filepath.Rel(clientDirectory, clientSpec.GoFileName)
	if err != nil {
		clientFilePath = clientSpec.GoFileName
	}

	structFilePath, err := filepath.Rel(
		clientDirectory,
		clientSpec.GoStructsFileName,
	)
	if err != nil {
		structFilePath = clientSpec.GoStructsFileName
	}

	// Return the client files
	return map[string][]byte{
		clientFilePath: client,
		structFilePath: structs,
	}, nil
}

/*
 * TChannel Client Generator
 */

// TCahnnelClientGenerator generates an instance of a zanzibar TChannel client
type TCahnnelClientGenerator struct {
	templates     *Template
	packageHelper *PackageHelper
}

// Generate returns the TChannel client generated files as a map of relative file
// path (relative to the target build directory) to file bytes.
func (generator *TCahnnelClientGenerator) Generate(
	instance *module.Instance,
) (map[string][]byte, error) {
	// Parse the client config from the endpoint JSON file
	clientConfig, err := readClientConfig(instance.JSONFileRaw)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading TChannel client %s JSON config",
			instance.InstanceName,
		)
	}

	clientSpec, err := NewTChannelClientSpec(
		filepath.Join(
			instance.BaseDirectory,
			instance.Directory,
			instance.JSONFileName,
		),
		clientConfig,
		generator.packageHelper,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error initializing TChannelClientSpec for %s",
			instance.InstanceName,
		)
	}

	clientMeta := &ClientMeta{
		PackageName:      clientSpec.ModuleSpec.PackageName,
		Services:         clientSpec.ModuleSpec.Services,
		IncludedPackages: clientSpec.ModuleSpec.IncludedPackages,
		ClientID:         clientSpec.ClientID,
	}

	client, err := generator.templates.execTemplate(
		"tchannel_client.tmpl",
		clientMeta,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel client template for %s",
			instance.InstanceName,
		)
	}

	server, err := generator.templates.execTemplate(
		"tchannel_server.tmpl",
		clientMeta,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel server template for %s",
			instance.InstanceName,
		)
	}

	handler, err := generator.templates.execTemplate(
		"tchannel_handler.tmpl",
		clientMeta,
	)
	if err != nil {
		return nil, errors.Wrapf(
			err,
			"Error executing TChannel handler template for %s",
			instance.InstanceName,
		)
	}

	clientDirectory := filepath.Join(
		generator.packageHelper.CodeGenTargetPath(),
		instance.Directory,
	)

	clientFilePath, err := filepath.Rel(clientDirectory, clientSpec.GoFileName)
	if err != nil {
		clientFilePath = clientSpec.GoFileName
	}

	// TODO:(lu) the locations of tchannel server and handler files are tentative
	serverFilePath := strings.TrimRight(clientFilePath, ".go") + "_server.go"
	handlerFilePath := strings.TrimRight(clientFilePath, ".go") + "_handler.go"

	// Return the client files
	return map[string][]byte{
		clientFilePath:  client,
		serverFilePath:  server,
		handlerFilePath: handler,
	}, nil
}

func readClientConfig(rawConfig []byte) (*clientClassConfig, error) {
	var clientConfig clientClassConfig
	if err := json.Unmarshal(rawConfig, &clientConfig); err != nil {
		return nil, errors.Wrapf(
			err,
			"Error reading config for HTTP client instance",
		)
	}
	clientConfig.Config["clientId"] = clientConfig.Name
	clientConfig.Config["clientType"] = clientConfig.Type
	return &clientConfig, nil
}
<|endoftext|>


<commit_before>package coinbasepro

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
	"github.com/thrasher-/gocryptotrader/common"
	"github.com/thrasher-/gocryptotrader/currency"
	exchange "github.com/thrasher-/gocryptotrader/exchanges"
	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
)

const (
	coinbaseproWebsocketURL = "wss://ws-feed.pro.coinbase.com"
)

// WebsocketSubscriber subscribes to websocket channels with respect to enabled
// currencies
func (c *CoinbasePro) WebsocketSubscriber() error {
	currencies := []string{}
	for _, x := range c.EnabledPairs.Strings() {
		currency := x[0:3] + "-" + x[3:]
		currencies = append(currencies, currency)
	}

	var channels = []WsChannels{
		{
			Name:       "heartbeat",
			ProductIDs: currencies,
		},
		{
			Name:       "ticker",
			ProductIDs: currencies,
		},
		{
			Name:       "level2",
			ProductIDs: currencies,
		},
	}

	subscribe := WebsocketSubscribe{Type: "subscribe", Channels: channels}

	data, err := common.JSONEncode(subscribe)
	if err != nil {
		return err
	}

	return c.WebsocketConn.WriteMessage(websocket.TextMessage, data)
}

// WsConnect initiates a websocket connection
func (c *CoinbasePro) WsConnect() error {
	if !c.Websocket.IsEnabled() || !c.IsEnabled() {
		return errors.New(exchange.WebsocketNotEnabled)
	}

	var dialer websocket.Dialer

	if c.Websocket.GetProxyAddress() != "" {
		proxy, err := url.Parse(c.Websocket.GetProxyAddress())
		if err != nil {
			return fmt.Errorf("coinbasepro_websocket.go error - proxy address %s",
				err)
		}

		dialer.Proxy = http.ProxyURL(proxy)
	}

	var err error
	c.WebsocketConn, _, err = dialer.Dial(c.Websocket.GetWebsocketURL(),
		http.Header{})
	if err != nil {
		return fmt.Errorf("coinbasepro_websocket.go error - unable to connect to websocket %s",
			err)
	}

	err = c.WebsocketSubscriber()
	if err != nil {
		return err
	}

	go c.WsHandleData()

	return nil
}

// WsReadData reads data from the websocket connection
func (c *CoinbasePro) WsReadData() (exchange.WebsocketResponse, error) {
	_, resp, err := c.WebsocketConn.ReadMessage()
	if err != nil {
		return exchange.WebsocketResponse{}, err
	}

	c.Websocket.TrafficAlert <- struct{}{}
	return exchange.WebsocketResponse{Raw: resp}, nil
}

// WsHandleData handles read data from websocket connection
func (c *CoinbasePro) WsHandleData() {
	c.Websocket.Wg.Add(1)

	defer func() {
		err := c.WebsocketConn.Close()
		if err != nil {
			c.Websocket.DataHandler <- fmt.Errorf("coinbasepro_websocket.go - Unable to to close Websocket connection. Error: %s",
				err)
		}
		c.Websocket.Wg.Done()
	}()

	for {
		select {
		case <-c.Websocket.ShutdownC:
			return

		default:
			resp, err := c.WsReadData()
			if err != nil {
				c.Websocket.DataHandler <- err
				return
			}

			type MsgType struct {
				Type      string `json:"type"`
				Sequence  int64  `json:"sequence"`
				ProductID string `json:"product_id"`
			}

			msgType := MsgType{}
			err = common.JSONDecode(resp.Raw, &msgType)
			if err != nil {
				c.Websocket.DataHandler <- err
				continue
			}

			if msgType.Type == "subscriptions" || msgType.Type == "heartbeat" {
				continue
			}

			switch msgType.Type {
			case "error":
				c.Websocket.DataHandler <- errors.New(string(resp.Raw))

			case "ticker":
				ticker := WebsocketTicker{}
				err := common.JSONDecode(resp.Raw, &ticker)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}

				c.Websocket.DataHandler <- exchange.TickerData{
					Timestamp: time.Now(),
					Pair:      currency.NewPairFromString(ticker.ProductID),
					AssetType: "SPOT",
					Exchange:  c.GetName(),
					OpenPrice: ticker.Price,
					HighPrice: ticker.High24H,
					LowPrice:  ticker.Low24H,
					Quantity:  ticker.Volume24H,
				}

			case "snapshot":
				snapshot := WebsocketOrderbookSnapshot{}
				err := common.JSONDecode(resp.Raw, &snapshot)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}

				err = c.ProcessSnapshot(&snapshot)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}

			case "l2update":
				update := WebsocketL2Update{}
				err := common.JSONDecode(resp.Raw, &update)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}

				err = c.ProcessUpdate(update)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}
			}
		}
	}
}

// ProcessSnapshot processes the initial orderbook snap shot
func (c *CoinbasePro) ProcessSnapshot(snapshot *WebsocketOrderbookSnapshot) error {
	var base *orderbook.Base
	for _, bid := range snapshot.Bids {
		price, err := strconv.ParseFloat(bid[0].(string), 64)
		if err != nil {
			return err
		}

		amount, err := strconv.ParseFloat(bid[1].(string), 64)
		if err != nil {
			return err
		}

		base.Bids = append(base.Bids,
			orderbook.Item{Price: price, Amount: amount})
	}

	for _, ask := range snapshot.Asks {
		price, err := strconv.ParseFloat(ask[0].(string), 64)
		if err != nil {
			return err
		}

		amount, err := strconv.ParseFloat(ask[1].(string), 64)
		if err != nil {
			return err
		}

		base.Asks = append(base.Asks,
			orderbook.Item{Price: price, Amount: amount})
	}

	pair := currency.NewPairFromString(snapshot.ProductID)
	base.AssetType = "SPOT"
	base.Pair = pair

	err := c.Websocket.Orderbook.LoadSnapshot(base, c.GetName(), false)
	if err != nil {
		return err
	}

	c.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
		Pair:     pair,
		Asset:    "SPOT",
		Exchange: c.GetName(),
	}

	return nil
}

// ProcessUpdate updates the orderbook local cache
func (c *CoinbasePro) ProcessUpdate(update WebsocketL2Update) error {
	var Asks, Bids []orderbook.Item

	for _, data := range update.Changes {
		price, _ := strconv.ParseFloat(data[1].(string), 64)
		volume, _ := strconv.ParseFloat(data[2].(string), 64)

		if data[0].(string) == "buy" {
			Bids = append(Bids, orderbook.Item{Price: price, Amount: volume})
		} else {
			Asks = append(Asks, orderbook.Item{Price: price, Amount: volume})
		}
	}

	if len(Asks) == 0 && len(Bids) == 0 {
		return errors.New("coibasepro_websocket.go error - no data in websocket update")
	}

	p := currency.NewPairFromString(update.ProductID)

	err := c.Websocket.Orderbook.Update(Bids, Asks, p, time.Now(), c.GetName(), "SPOT")
	if err != nil {
		return err
	}

	c.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
		Pair:     p,
		Asset:    "SPOT",
		Exchange: c.GetName(),
	}

	return nil
}
<commit_msg>Fix CoinbasePro websocket bug introduced in PR #262 (#269)<commit_after>package coinbasepro

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
	"github.com/thrasher-/gocryptotrader/common"
	"github.com/thrasher-/gocryptotrader/currency"
	exchange "github.com/thrasher-/gocryptotrader/exchanges"
	"github.com/thrasher-/gocryptotrader/exchanges/orderbook"
)

const (
	coinbaseproWebsocketURL = "wss://ws-feed.pro.coinbase.com"
)

// WebsocketSubscriber subscribes to websocket channels with respect to enabled
// currencies
func (c *CoinbasePro) WebsocketSubscriber() error {
	currencies := []string{}
	for _, x := range c.EnabledPairs.Strings() {
		currency := x[0:3] + "-" + x[3:]
		currencies = append(currencies, currency)
	}

	var channels = []WsChannels{
		{
			Name:       "heartbeat",
			ProductIDs: currencies,
		},
		{
			Name:       "ticker",
			ProductIDs: currencies,
		},
		{
			Name:       "level2",
			ProductIDs: currencies,
		},
	}

	subscribe := WebsocketSubscribe{Type: "subscribe", Channels: channels}

	data, err := common.JSONEncode(subscribe)
	if err != nil {
		return err
	}

	return c.WebsocketConn.WriteMessage(websocket.TextMessage, data)
}

// WsConnect initiates a websocket connection
func (c *CoinbasePro) WsConnect() error {
	if !c.Websocket.IsEnabled() || !c.IsEnabled() {
		return errors.New(exchange.WebsocketNotEnabled)
	}

	var dialer websocket.Dialer

	if c.Websocket.GetProxyAddress() != "" {
		proxy, err := url.Parse(c.Websocket.GetProxyAddress())
		if err != nil {
			return fmt.Errorf("coinbasepro_websocket.go error - proxy address %s",
				err)
		}

		dialer.Proxy = http.ProxyURL(proxy)
	}

	var err error
	c.WebsocketConn, _, err = dialer.Dial(c.Websocket.GetWebsocketURL(),
		http.Header{})
	if err != nil {
		return fmt.Errorf("coinbasepro_websocket.go error - unable to connect to websocket %s",
			err)
	}

	err = c.WebsocketSubscriber()
	if err != nil {
		return err
	}

	go c.WsHandleData()

	return nil
}

// WsReadData reads data from the websocket connection
func (c *CoinbasePro) WsReadData() (exchange.WebsocketResponse, error) {
	_, resp, err := c.WebsocketConn.ReadMessage()
	if err != nil {
		return exchange.WebsocketResponse{}, err
	}

	c.Websocket.TrafficAlert <- struct{}{}
	return exchange.WebsocketResponse{Raw: resp}, nil
}

// WsHandleData handles read data from websocket connection
func (c *CoinbasePro) WsHandleData() {
	c.Websocket.Wg.Add(1)

	defer func() {
		err := c.WebsocketConn.Close()
		if err != nil {
			c.Websocket.DataHandler <- fmt.Errorf("coinbasepro_websocket.go - Unable to to close Websocket connection. Error: %s",
				err)
		}
		c.Websocket.Wg.Done()
	}()

	for {
		select {
		case <-c.Websocket.ShutdownC:
			return

		default:
			resp, err := c.WsReadData()
			if err != nil {
				c.Websocket.DataHandler <- err
				return
			}

			type MsgType struct {
				Type      string `json:"type"`
				Sequence  int64  `json:"sequence"`
				ProductID string `json:"product_id"`
			}

			msgType := MsgType{}
			err = common.JSONDecode(resp.Raw, &msgType)
			if err != nil {
				c.Websocket.DataHandler <- err
				continue
			}

			if msgType.Type == "subscriptions" || msgType.Type == "heartbeat" {
				continue
			}

			switch msgType.Type {
			case "error":
				c.Websocket.DataHandler <- errors.New(string(resp.Raw))

			case "ticker":
				ticker := WebsocketTicker{}
				err := common.JSONDecode(resp.Raw, &ticker)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}

				c.Websocket.DataHandler <- exchange.TickerData{
					Timestamp: time.Now(),
					Pair:      currency.NewPairFromString(ticker.ProductID),
					AssetType: "SPOT",
					Exchange:  c.GetName(),
					OpenPrice: ticker.Price,
					HighPrice: ticker.High24H,
					LowPrice:  ticker.Low24H,
					Quantity:  ticker.Volume24H,
				}

			case "snapshot":
				snapshot := WebsocketOrderbookSnapshot{}
				err := common.JSONDecode(resp.Raw, &snapshot)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}

				err = c.ProcessSnapshot(&snapshot)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}

			case "l2update":
				update := WebsocketL2Update{}
				err := common.JSONDecode(resp.Raw, &update)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}

				err = c.ProcessUpdate(update)
				if err != nil {
					c.Websocket.DataHandler <- err
					continue
				}
			}
		}
	}
}

// ProcessSnapshot processes the initial orderbook snap shot
func (c *CoinbasePro) ProcessSnapshot(snapshot *WebsocketOrderbookSnapshot) error {
	var base orderbook.Base
	for _, bid := range snapshot.Bids {
		price, err := strconv.ParseFloat(bid[0].(string), 64)
		if err != nil {
			return err
		}

		amount, err := strconv.ParseFloat(bid[1].(string), 64)
		if err != nil {
			return err
		}

		base.Bids = append(base.Bids,
			orderbook.Item{Price: price, Amount: amount})
	}

	for _, ask := range snapshot.Asks {
		price, err := strconv.ParseFloat(ask[0].(string), 64)
		if err != nil {
			return err
		}

		amount, err := strconv.ParseFloat(ask[1].(string), 64)
		if err != nil {
			return err
		}

		base.Asks = append(base.Asks,
			orderbook.Item{Price: price, Amount: amount})
	}

	pair := currency.NewPairFromString(snapshot.ProductID)
	base.AssetType = "SPOT"
	base.Pair = pair

	err := c.Websocket.Orderbook.LoadSnapshot(&base, c.GetName(), false)
	if err != nil {
		return err
	}

	c.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
		Pair:     pair,
		Asset:    "SPOT",
		Exchange: c.GetName(),
	}

	return nil
}

// ProcessUpdate updates the orderbook local cache
func (c *CoinbasePro) ProcessUpdate(update WebsocketL2Update) error {
	var Asks, Bids []orderbook.Item

	for _, data := range update.Changes {
		price, _ := strconv.ParseFloat(data[1].(string), 64)
		volume, _ := strconv.ParseFloat(data[2].(string), 64)

		if data[0].(string) == "buy" {
			Bids = append(Bids, orderbook.Item{Price: price, Amount: volume})
		} else {
			Asks = append(Asks, orderbook.Item{Price: price, Amount: volume})
		}
	}

	if len(Asks) == 0 && len(Bids) == 0 {
		return errors.New("coibasepro_websocket.go error - no data in websocket update")
	}

	p := currency.NewPairFromString(update.ProductID)

	err := c.Websocket.Orderbook.Update(Bids, Asks, p, time.Now(), c.GetName(), "SPOT")
	if err != nil {
		return err
	}

	c.Websocket.DataHandler <- exchange.WebsocketOrderbookUpdate{
		Pair:     p,
		Asset:    "SPOT",
		Exchange: c.GetName(),
	}

	return nil
}
<|endoftext|>


<commit_before>package fastq

import (
	"bytes"
	"fmt"
	"gongs/scan"
	"gongs/xopen"
	"io"
	"strings"
)

type Fastq struct {
	Name string
	Seq  []byte
	Qual []byte
}

func (fq Fastq) String() string {
	return fmt.Sprintf("@%v\n%v\n+\n%v", fq.Name, string(fq.Seq), string(fq.Qual))
}

func (fq Fastq) Id() string {
	if n := strings.IndexByte(fq.Name, ' '); n >= 0 {
		return fq.Name[:n]
	}

	// for old solexa data format
	if n := strings.IndexByte(fq.Name, '#'); n >= 0 {
		return fq.Name[:n]
	}

	return fq.Name
}

type FastqFile struct {
	Name  string
	file  io.ReadCloser
	s     *scan.Scanner
	name  []byte
	seq   []byte
	qual  []byte
	err   error
	stage int
}

func Open(filename string) (*FastqFile, error) {
	file, err := xopen.Xopen(filename)
	if err != nil {
		return nil, err
	}

	return &FastqFile{
		Name: filename,
		s:    scan.New(file),
		file: file,
	}, nil
}

func (ff *FastqFile) Close() error {
	return ff.file.Close()
}

func (ff *FastqFile) Err() error {
	if ff.err == nil || ff.err == io.EOF {
		if err := ff.s.Err(); err != nil {
			return err
		}
		return nil
	}
	return ff.err
}

func (ff *FastqFile) setErr(err error) {
	if ff.err == nil {
		ff.err = err
	}
}

func (ff *FastqFile) Next() bool {
	if ff.err != nil {
		return false
	}

	var line []byte
	for ff.s.Scan() {
		line = bytes.TrimSpace(ff.s.Bytes())
		if len(line) == 0 { // ingore empty line
			continue
		}
		switch ff.stage {
		case 0: // get fastq name
			if len(line) > 0 && line[0] != '@' {
				ff.setErr(fmt.Errorf("file: %v Wrong Fastq Record Name %s at line: %d", ff.Name, string(line), ff.s.Lid()))
				return false
			}
			ff.stage++
			ff.name = line[1:]
			ff.seq = ff.seq[:0]   // clear seq
			ff.qual = ff.qual[:0] // clear qual
		case 1: // get fastq seq
			if len(line) > 0 && line[0] == '+' {
				ff.stage += 2
				break
			}
			ff.seq = append(ff.seq, line...)
		case 2: // get + line
		case 3: // get fastq qual
			ff.qual = append(ff.qual, line...)
			if len(ff.qual) == len(ff.seq) {
				ff.stage = 0
				return true
			} else if len(ff.qual) > len(ff.seq) {
				ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual length (%d) != seq length (%d) at line: %d",
					ff.Name, string(ff.name), len(ff.qual), len(ff.seq), ff.s.Lid()))
				return false
			}
		}
	}
	if len(ff.qual) < len(ff.seq) {
		ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual length (%d) != seq length (%d) at line: %d",
			ff.Name, string(ff.name), len(ff.qual), len(ff.seq), ff.s.Lid()))
	}
	ff.setErr(io.EOF)
	return false
}

func (ff *FastqFile) Value() *Fastq {
	return &Fastq{Name: string(ff.name), Seq: ff.seq, Qual: ff.qual}
}

func (ff *FastqFile) Iter() <-chan *Fastq {
	ch := make(chan *Fastq)
	go func(ch chan *Fastq) {
		for ff.Next() {
			ch <- ff.Value()
		}
		close(ch)
	}(ch)
	return ch
}
<commit_msg>add Opens function<commit_after>package fastq

import (
	"biofile/fastq"
	"bytes"
	"fmt"
	"gongs/scan"
	"gongs/xopen"
	"io"
	"strings"
)

type Fastq struct {
	Name string
	Seq  []byte
	Qual []byte
}

func (fq Fastq) String() string {
	return fmt.Sprintf("@%v\n%v\n+\n%v", fq.Name, string(fq.Seq), string(fq.Qual))
}

func (fq Fastq) Id() string {
	if n := strings.IndexByte(fq.Name, ' '); n >= 0 {
		return fq.Name[:n]
	}

	// for old solexa data format
	if n := strings.IndexByte(fq.Name, '#'); n >= 0 {
		return fq.Name[:n]
	}

	return fq.Name
}

type FastqFile struct {
	Name  string
	file  io.ReadCloser
	s     *scan.Scanner
	name  []byte
	seq   []byte
	qual  []byte
	err   error
	stage int
}

func Open(filename string) (*FastqFile, error) {
	file, err := xopen.Xopen(filename)
	if err != nil {
		return nil, err
	}

	return &FastqFile{
		Name: filename,
		s:    scan.New(file),
		file: file,
	}, nil
}

func (ff *FastqFile) Close() error {
	return ff.file.Close()
}

func (ff *FastqFile) Err() error {
	if ff.err == nil || ff.err == io.EOF {
		if err := ff.s.Err(); err != nil {
			return err
		}
		return nil
	}
	return ff.err
}

func (ff *FastqFile) setErr(err error) {
	if ff.err == nil {
		ff.err = err
	}
}

func (ff *FastqFile) Next() bool {
	if ff.err != nil {
		return false
	}

	var line []byte
	for ff.s.Scan() {
		line = bytes.TrimSpace(ff.s.Bytes())
		if len(line) == 0 { // ingore empty line
			continue
		}
		switch ff.stage {
		case 0: // get fastq name
			if len(line) > 0 && line[0] != '@' {
				ff.setErr(fmt.Errorf("file: %v Wrong Fastq Record Name %s at line: %d", ff.Name, string(line), ff.s.Lid()))
				return false
			}
			ff.stage++
			ff.name = line[1:]
			ff.seq = ff.seq[:0]   // clear seq
			ff.qual = ff.qual[:0] // clear qual
		case 1: // get fastq seq
			if len(line) > 0 && line[0] == '+' {
				ff.stage += 2
				break
			}
			ff.seq = append(ff.seq, line...)
		case 2: // get + line
		case 3: // get fastq qual
			ff.qual = append(ff.qual, line...)
			if len(ff.qual) == len(ff.seq) {
				ff.stage = 0
				return true
			} else if len(ff.qual) > len(ff.seq) {
				ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual length (%d) != seq length (%d) at line: %d",
					ff.Name, string(ff.name), len(ff.qual), len(ff.seq), ff.s.Lid()))
				return false
			}
		}
	}
	if len(ff.qual) < len(ff.seq) {
		ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual length (%d) != seq length (%d) at line: %d",
			ff.Name, string(ff.name), len(ff.qual), len(ff.seq), ff.s.Lid()))
	}
	ff.setErr(io.EOF)
	return false
}

func (ff *FastqFile) Value() *Fastq {
	return &Fastq{Name: string(ff.name), Seq: ff.seq, Qual: ff.qual}
}

func (ff *FastqFile) Iter() <-chan *Fastq {
	ch := make(chan *Fastq)
	go func(ch chan *Fastq) {
		for ff.Next() {
			ch <- ff.Value()
		}
		close(ch)
	}(ch)
	return ch
}

func Opens(filenames ...string) ([]*FastqFile, error) {
	fqfiles := make([]*fastq.FastqFile, len(filenames))
	for i, filename := range filenames {
		fqfile, err := fastq.Open(filename)
		if err != nil {
			return err
		}
		fqfiles[i] = fqfile
	}
	return fqfiles, nil
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"time"

	"github.com/cncd/logging"
	"github.com/cncd/pubsub"
	"github.com/drone/drone/model"
	"github.com/drone/drone/router/middleware/session"
	"github.com/drone/drone/store"

	"github.com/Sirupsen/logrus"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

var (
	// Time allowed to write the file to the client.
	writeWait = 5 * time.Second

	// Time allowed to read the next pong message from the client.
	pongWait = 60 * time.Second

	// Send pings to client with this period. Must be less than pongWait.
	pingPeriod = 30 * time.Second

	// upgrader defines the default behavior for upgrading the websocket.
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
)

func reader(ws *websocket.Conn) {
	defer ws.Close()
	ws.SetReadLimit(512)
	ws.SetReadDeadline(time.Now().Add(pongWait))
	ws.SetPongHandler(func(string) error {
		ws.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})
	for {
		_, _, err := ws.ReadMessage()
		if err != nil {
			break
		}
	}
}

func LogStream(c *gin.Context) {
	repo := session.Repo(c)
	buildn, _ := strconv.Atoi(c.Param("build"))
	jobn, _ := strconv.Atoi(c.Param("number"))

	build, err := store.GetBuildNumber(c, repo, buildn)
	if err != nil {
		logrus.Debugln("stream cannot get build number.", err)
		c.AbortWithError(404, err)
		return
	}
	proc, err := store.FromContext(c).ProcFind(build, jobn)
	if err != nil {
		logrus.Debugln("stream cannot get proc number.", err)
		c.AbortWithError(404, err)
		return
	}
	if proc.State != model.StatusRunning {
		logrus.Debugln("stream not found.")
		c.AbortWithStatus(404)
		return
	}

	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		if _, ok := err.(websocket.HandshakeError); !ok {
			logrus.Errorf("Cannot upgrade websocket. %s", err)
		}
		return
	}
	logrus.Debugf("Successfull upgraded websocket")

	ticker := time.NewTicker(pingPeriod)
	logc := make(chan []byte, 10)

	ctx, cancel := context.WithCancel(
		context.Background(),
	)
	defer func() {
		cancel()
		ticker.Stop()
		close(logc)
		logrus.Debugf("Successfully closing websocket")
	}()

	go func() {
		// TODO remove global variable
		Config.Services.Logs.Tail(ctx, fmt.Sprint(proc.ID), func(entries ...*logging.Entry) {
			for _, entry := range entries {
				select {
				case <-ctx.Done():
					return
				default:
					logc <- entry.Data
				}
			}
		})
		cancel()
	}()

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case buf, ok := <-logc:
				if ok {
					ws.SetWriteDeadline(time.Now().Add(writeWait))
					ws.WriteMessage(websocket.TextMessage, buf)
				}
			case <-ticker.C:
				err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait))
				if err != nil {
					return
				}
			}
		}
	}()
	reader(ws)
}

func EventStream(c *gin.Context) {
	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		if _, ok := err.(websocket.HandshakeError); !ok {
			logrus.Errorf("Cannot upgrade websocket. %s", err)
		}
		return
	}
	logrus.Debugf("Successfull upgraded websocket")

	user := session.User(c)
	repo := map[string]bool{}
	if user != nil {
		repos, _ := store.FromContext(c).RepoList(user)
		for _, r := range repos {
			repo[r.FullName] = true
		}
	}

	ticker := time.NewTicker(pingPeriod)
	eventc := make(chan []byte, 10)

	ctx, cancel := context.WithCancel(
		context.Background(),
	)
	defer func() {
		cancel()
		ticker.Stop()
		close(eventc)
		logrus.Debugf("Successfully closing websocket")
	}()

	go func() {
		// TODO remove this from global config
		Config.Services.Pubsub.Subscribe(c, "topic/events", func(m pubsub.Message) {
			name := m.Labels["repo"]
			priv := m.Labels["private"]
			if repo[name] || priv == "false" {
				select {
				case <-ctx.Done():
					return
				default:
					eventc <- m.Data
				}
			}
		})
		cancel()
	}()

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case buf, ok := <-eventc:
				if ok {
					ws.SetWriteDeadline(time.Now().Add(writeWait))
					ws.WriteMessage(websocket.TextMessage, buf)
				}
			case <-ticker.C:
				err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait))
				if err != nil {
					return
				}
			}
		}
	}()
	reader(ws)
}

//
// event source streaming for compatibility with quic and http2
//

func EventStreamSSE(c *gin.Context) {
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no")

	rw := c.Writer

	flusher, ok := rw.(http.Flusher)
	if !ok {
		c.String(500, "Streaming not supported")
		return
	}

	logrus.Debugf("user feed: connection opened")

	user := session.User(c)
	repo := map[string]bool{}
	if user != nil {
		repos, _ := store.FromContext(c).RepoList(user)
		for _, r := range repos {
			repo[r.FullName] = true
		}
	}

	eventc := make(chan []byte, 10)
	ctx, cancel := context.WithCancel(
		context.Background(),
	)

	defer func() {
		cancel()
		close(eventc)
		logrus.Debugf("user feed: connection closed")
	}()

	go func() {
		// TODO remove this from global config
		Config.Services.Pubsub.Subscribe(c, "topic/events", func(m pubsub.Message) {
			name := m.Labels["repo"]
			priv := m.Labels["private"]
			if repo[name] || priv == "false" {
				select {
				case <-ctx.Done():
					return
				default:
					eventc <- m.Data
				}
			}
		})
		cancel()
	}()

	for {
		select {
		case <-rw.CloseNotify():
			return
		case <-ctx.Done():
			return
		case buf, ok := <-eventc:
			if ok {
				io.WriteString(rw, "data: ")
				rw.Write(buf)
				io.WriteString(rw, "\n\n")
				flusher.Flush()
			}
		}
	}
}

func LogStreamSSE(c *gin.Context) {
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no")

	rw := c.Writer

	flusher, ok := rw.(http.Flusher)
	if !ok {
		c.String(500, "Streaming not supported")
		return
	}

	// repo := session.Repo(c)
	//
	// // parse the build number and job sequence number from
	// // the repquest parameter.
	// num, _ := strconv.Atoi(c.Params.ByName("number"))
	// ppid, _ := strconv.Atoi(c.Params.ByName("ppid"))
	// name := c.Params.ByName("proc")
	//
	// build, err := store.GetBuildNumber(c, repo, num)
	// if err != nil {
	// 	c.AbortWithError(404, err)
	// 	return
	// }
	//
	// proc, err := store.FromContext(c).ProcChild(build, ppid, name)
	// if err != nil {
	// 	c.AbortWithError(404, err)
	// 	return
	// }

	repo := session.Repo(c)
	buildn, _ := strconv.Atoi(c.Param("build"))
	jobn, _ := strconv.Atoi(c.Param("number"))

	build, err := store.GetBuildNumber(c, repo, buildn)
	if err != nil {
		logrus.Debugln("stream cannot get build number.", err)
		io.WriteString(rw, "event: error\ndata: build not found\n\n")
		return
	}
	proc, err := store.FromContext(c).ProcFind(build, jobn)
	if err != nil {
		logrus.Debugln("stream cannot get proc number.", err)
		io.WriteString(rw, "event: error\ndata: process not found\n\n")
		return
	}
	if proc.State != model.StatusRunning {
		logrus.Debugln("stream not found.")
		io.WriteString(rw, "event: error\ndata: stream not found\n\n")
		return
	}

	logc := make(chan []byte, 10)
	ctx, cancel := context.WithCancel(
		context.Background(),
	)

	logrus.Debugf("log stream: connection opened")

	defer func() {
		cancel()
		close(logc)
		logrus.Debugf("log stream: connection closed")
	}()

	go func() {
		// TODO remove global variable
		Config.Services.Logs.Tail(ctx, fmt.Sprint(proc.ID), func(entries ...*logging.Entry) {
			for _, entry := range entries {
				select {
				case <-ctx.Done():
					return
				default:
					logc <- entry.Data
				}
			}
		})

		io.WriteString(rw, "event: error\ndata: eof\n\n")

		cancel()
	}()

	id := 1
	last, _ := strconv.Atoi(
		c.Request.Header.Get("Last-Event-ID"),
	)
	if last != 0 {
		logrus.Debugf("log stream: reconnect: last-event-id: %d", last)
	}

	// retry: 10000\n

	for {
		select {
		// after 1 hour of idle (no response) end the stream.
		// this is more of a safety mechanism than anything,
		// and can be removed once the code is more mature.
		case <-time.After(time.Hour):
			return
		case <-rw.CloseNotify():
			return
		case <-ctx.Done():
			return
		case buf, ok := <-logc:
			if ok {
				if id > last {
					io.WriteString(rw, "id: "+strconv.Itoa(id))
					io.WriteString(rw, "\n")
					io.WriteString(rw, "data: ")
					rw.Write(buf)
					io.WriteString(rw, "\n\n")
					flusher.Flush()
				}
				id++
			}
		}
	}
}
<commit_msg>periodically ping client from server<commit_after>package server

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"time"

	"github.com/cncd/logging"
	"github.com/cncd/pubsub"
	"github.com/drone/drone/model"
	"github.com/drone/drone/router/middleware/session"
	"github.com/drone/drone/store"

	"github.com/Sirupsen/logrus"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

var (
	// Time allowed to write the file to the client.
	writeWait = 5 * time.Second

	// Time allowed to read the next pong message from the client.
	pongWait = 60 * time.Second

	// Send pings to client with this period. Must be less than pongWait.
	pingPeriod = 30 * time.Second

	// upgrader defines the default behavior for upgrading the websocket.
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
)

func reader(ws *websocket.Conn) {
	defer ws.Close()
	ws.SetReadLimit(512)
	ws.SetReadDeadline(time.Now().Add(pongWait))
	ws.SetPongHandler(func(string) error {
		ws.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})
	for {
		_, _, err := ws.ReadMessage()
		if err != nil {
			break
		}
	}
}

func LogStream(c *gin.Context) {
	repo := session.Repo(c)
	buildn, _ := strconv.Atoi(c.Param("build"))
	jobn, _ := strconv.Atoi(c.Param("number"))

	build, err := store.GetBuildNumber(c, repo, buildn)
	if err != nil {
		logrus.Debugln("stream cannot get build number.", err)
		c.AbortWithError(404, err)
		return
	}
	proc, err := store.FromContext(c).ProcFind(build, jobn)
	if err != nil {
		logrus.Debugln("stream cannot get proc number.", err)
		c.AbortWithError(404, err)
		return
	}
	if proc.State != model.StatusRunning {
		logrus.Debugln("stream not found.")
		c.AbortWithStatus(404)
		return
	}

	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		if _, ok := err.(websocket.HandshakeError); !ok {
			logrus.Errorf("Cannot upgrade websocket. %s", err)
		}
		return
	}
	logrus.Debugf("Successfull upgraded websocket")

	ticker := time.NewTicker(pingPeriod)
	logc := make(chan []byte, 10)

	ctx, cancel := context.WithCancel(
		context.Background(),
	)
	defer func() {
		cancel()
		ticker.Stop()
		close(logc)
		logrus.Debugf("Successfully closing websocket")
	}()

	go func() {
		// TODO remove global variable
		Config.Services.Logs.Tail(ctx, fmt.Sprint(proc.ID), func(entries ...*logging.Entry) {
			for _, entry := range entries {
				select {
				case <-ctx.Done():
					return
				default:
					logc <- entry.Data
				}
			}
		})
		cancel()
	}()

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case buf, ok := <-logc:
				if ok {
					ws.SetWriteDeadline(time.Now().Add(writeWait))
					ws.WriteMessage(websocket.TextMessage, buf)
				}
			case <-ticker.C:
				err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait))
				if err != nil {
					return
				}
			}
		}
	}()
	reader(ws)
}

func EventStream(c *gin.Context) {
	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		if _, ok := err.(websocket.HandshakeError); !ok {
			logrus.Errorf("Cannot upgrade websocket. %s", err)
		}
		return
	}
	logrus.Debugf("Successfull upgraded websocket")

	user := session.User(c)
	repo := map[string]bool{}
	if user != nil {
		repos, _ := store.FromContext(c).RepoList(user)
		for _, r := range repos {
			repo[r.FullName] = true
		}
	}

	ticker := time.NewTicker(pingPeriod)
	eventc := make(chan []byte, 10)

	ctx, cancel := context.WithCancel(
		context.Background(),
	)
	defer func() {
		cancel()
		ticker.Stop()
		close(eventc)
		logrus.Debugf("Successfully closing websocket")
	}()

	go func() {
		// TODO remove this from global config
		Config.Services.Pubsub.Subscribe(c, "topic/events", func(m pubsub.Message) {
			name := m.Labels["repo"]
			priv := m.Labels["private"]
			if repo[name] || priv == "false" {
				select {
				case <-ctx.Done():
					return
				default:
					eventc <- m.Data
				}
			}
		})
		cancel()
	}()

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case buf, ok := <-eventc:
				if ok {
					ws.SetWriteDeadline(time.Now().Add(writeWait))
					ws.WriteMessage(websocket.TextMessage, buf)
				}
			case <-ticker.C:
				err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait))
				if err != nil {
					return
				}
			}
		}
	}()
	reader(ws)
}

//
// event source streaming for compatibility with quic and http2
//

func EventStreamSSE(c *gin.Context) {
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no")

	rw := c.Writer

	flusher, ok := rw.(http.Flusher)
	if !ok {
		c.String(500, "Streaming not supported")
		return
	}

	// ping the client
	io.WriteString(rw, ": ping\n\n")
	flusher.Flush()

	logrus.Debugf("user feed: connection opened")

	user := session.User(c)
	repo := map[string]bool{}
	if user != nil {
		repos, _ := store.FromContext(c).RepoList(user)
		for _, r := range repos {
			repo[r.FullName] = true
		}
	}

	eventc := make(chan []byte, 10)
	ctx, cancel := context.WithCancel(
		context.Background(),
	)

	defer func() {
		cancel()
		close(eventc)
		logrus.Debugf("user feed: connection closed")
	}()

	go func() {
		// TODO remove this from global config
		Config.Services.Pubsub.Subscribe(c, "topic/events", func(m pubsub.Message) {
			name := m.Labels["repo"]
			priv := m.Labels["private"]
			if repo[name] || priv == "false" {
				select {
				case <-ctx.Done():
					return
				default:
					eventc <- m.Data
				}
			}
		})
		cancel()
	}()

	for {
		select {
		case <-rw.CloseNotify():
			return
		case <-ctx.Done():
			return
		case <-time.After(time.Second * 30):
			io.WriteString(rw, ": ping\n\n")
			flusher.Flush()
		case buf, ok := <-eventc:
			if ok {
				io.WriteString(rw, "data: ")
				rw.Write(buf)
				io.WriteString(rw, "\n\n")
				flusher.Flush()
			}
		}
	}
}

func LogStreamSSE(c *gin.Context) {
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no")

	rw := c.Writer

	flusher, ok := rw.(http.Flusher)
	if !ok {
		c.String(500, "Streaming not supported")
		return
	}

	io.WriteString(rw, ": ping\n\n")
	flusher.Flush()

	// repo := session.Repo(c)
	//
	// // parse the build number and job sequence number from
	// // the repquest parameter.
	// num, _ := strconv.Atoi(c.Params.ByName("number"))
	// ppid, _ := strconv.Atoi(c.Params.ByName("ppid"))
	// name := c.Params.ByName("proc")
	//
	// build, err := store.GetBuildNumber(c, repo, num)
	// if err != nil {
	// 	c.AbortWithError(404, err)
	// 	return
	// }
	//
	// proc, err := store.FromContext(c).ProcChild(build, ppid, name)
	// if err != nil {
	// 	c.AbortWithError(404, err)
	// 	return
	// }

	repo := session.Repo(c)
	buildn, _ := strconv.Atoi(c.Param("build"))
	jobn, _ := strconv.Atoi(c.Param("number"))

	build, err := store.GetBuildNumber(c, repo, buildn)
	if err != nil {
		logrus.Debugln("stream cannot get build number.", err)
		io.WriteString(rw, "event: error\ndata: build not found\n\n")
		return
	}
	proc, err := store.FromContext(c).ProcFind(build, jobn)
	if err != nil {
		logrus.Debugln("stream cannot get proc number.", err)
		io.WriteString(rw, "event: error\ndata: process not found\n\n")
		return
	}
	if proc.State != model.StatusRunning {
		logrus.Debugln("stream not found.")
		io.WriteString(rw, "event: error\ndata: stream not found\n\n")
		return
	}

	logc := make(chan []byte, 10)
	ctx, cancel := context.WithCancel(
		context.Background(),
	)

	logrus.Debugf("log stream: connection opened")

	defer func() {
		cancel()
		close(logc)
		logrus.Debugf("log stream: connection closed")
	}()

	go func() {
		// TODO remove global variable
		Config.Services.Logs.Tail(ctx, fmt.Sprint(proc.ID), func(entries ...*logging.Entry) {
			for _, entry := range entries {
				select {
				case <-ctx.Done():
					return
				default:
					logc <- entry.Data
				}
			}
		})

		io.WriteString(rw, "event: error\ndata: eof\n\n")

		cancel()
	}()

	id := 1
	last, _ := strconv.Atoi(
		c.Request.Header.Get("Last-Event-ID"),
	)
	if last != 0 {
		logrus.Debugf("log stream: reconnect: last-event-id: %d", last)
	}

	// retry: 10000\n

	for {
		select {
		// after 1 hour of idle (no response) end the stream.
		// this is more of a safety mechanism than anything,
		// and can be removed once the code is more mature.
		case <-time.After(time.Hour):
			return
		case <-rw.CloseNotify():
			return
		case <-ctx.Done():
			return
		case <-time.After(time.Second * 30):
			io.WriteString(rw, ": ping\n\n")
			flusher.Flush()
		case buf, ok := <-logc:
			if ok {
				if id > last {
					io.WriteString(rw, "id: "+strconv.Itoa(id))
					io.WriteString(rw, "\n")
					io.WriteString(rw, "data: ")
					rw.Write(buf)
					io.WriteString(rw, "\n\n")
					flusher.Flush()
				}
				id++
			}
		}
	}
}
<|endoftext|>