<commit_before>// The client package helps developers connect to Gearmand, send
// jobs and fetch result.
package client

import (
	"bufio"
	"net"
	"sync"
)

// One client connect to one server.
// Use Pool for multi-connections.
type Client struct {
	sync.Mutex

	net, addr, lastcall string
	respHandler         map[string]ResponseHandler
	innerHandler        map[string]ResponseHandler
	in                  chan *Response
	conn                net.Conn
	rw                  *bufio.ReadWriter

	ErrorHandler ErrorHandler
}

// Return a client.
func New(network, addr string) (client *Client, err error) {
	client = &Client{
		net:          network,
		addr:         addr,
		respHandler:  make(map[string]ResponseHandler, queueSize),
		innerHandler: make(map[string]ResponseHandler, queueSize),
		in:           make(chan *Response, queueSize),
	}
	client.conn, err = net.Dial(client.net, client.addr)
	if err != nil {
		return
	}
	client.rw = bufio.NewReadWriter(bufio.NewReader(client.conn),
		bufio.NewWriter(client.conn))
	go client.readLoop()
	go client.processLoop()
	return
}

func (client *Client) write(req *request) (err error) {
	var n int
	buf := req.Encode()
	for i := 0; i < len(buf); i += n {
		n, err = client.rw.Write(buf[i:])
		if err != nil {
			return
		}
	}
	return client.rw.Flush()
}

func (client *Client) read(length int) (data []byte, err error) {
	n := 0
	buf := getBuffer(bufferSize)
	// read until data can be unpacked
	for i := length; i > 0 || len(data) < minPacketLength; i -= n {
		if n, err = client.rw.Read(buf); err != nil {
			return
		}
		data = append(data, buf[0:n]...)
		if n < bufferSize {
			break
		}
	}
	return
}

func (client *Client) readLoop() {
	defer close(client.in)
	var data, leftdata []byte
	var err error
	var resp *Response
ReadLoop:
	for client.conn != nil {
		if data, err = client.read(bufferSize); err != nil {
			if opErr, ok := err.(*net.OpError); ok {
				if opErr.Timeout() {
					client.err(err)
				}
				if opErr.Temporary() {
					continue
				}
				break
			}
			client.err(err)
			// If it is unexpected error and the connection wasn't
			// closed by Gearmand, the client should close the conection
			// and reconnect to job server.
			client.Close()
			client.conn, err = net.Dial(client.net, client.addr)
			if err != nil {
				client.err(err)
				break
			}
			client.rw = bufio.NewReadWriter(bufio.NewReader(client.conn),
				bufio.NewWriter(client.conn))
			continue
		}
		if len(leftdata) > 0 { // some data left for processing
			data = append(leftdata, data...)
			leftdata = nil
		}
		for {
			l := len(data)
			if l < minPacketLength { // not enough data
				leftdata = data
				continue ReadLoop
			}
			if resp, l, err = decodeResponse(data); err != nil {
				leftdata = data[l:]
				continue ReadLoop
			} else {
				client.in <- resp
			}
			data = data[l:]
			if len(data) > 0 {
				continue
			}
			break
		}
	}
}

func (client *Client) processLoop() {
	for resp := range client.in {
		switch resp.DataType {
		case dtError:
			if client.lastcall != "" {
				resp = client.handleInner(client.lastcall, resp)
				client.lastcall = ""
			} else {
				client.err(getError(resp.Data))
			}
		case dtStatusRes:
			resp = client.handleInner("s"+resp.Handle, resp)
		case dtJobCreated:
			resp = client.handleInner("c", resp)
		case dtEchoRes:
			resp = client.handleInner("e", resp)
		case dtWorkData, dtWorkWarning, dtWorkStatus:
			resp = client.handleResponse(resp.Handle, resp)
		case dtWorkComplete, dtWorkFail, dtWorkException:
			client.handleResponse(resp.Handle, resp)
			delete(client.respHandler, resp.Handle)
		}
	}
}

func (client *Client) err(e error) {
	if client.ErrorHandler != nil {
		client.ErrorHandler(e)
	}
}

func (client *Client) handleResponse(key string, resp *Response) *Response {
	if h, ok := client.respHandler[key]; ok {
		h(resp)
		return nil
	}
	return resp
}

func (client *Client) handleInner(key string, resp *Response) *Response {
	if h, ok := client.innerHandler[key]; ok {
		h(resp)
		delete(client.innerHandler, key)
		return nil
	}
	return resp
}

func (client *Client) do(funcname string, data []byte,
	flag uint32) (handle string, err error) {
	if client.conn == nil {
		return "", ErrLostConn
	}
	var mutex sync.Mutex
	mutex.Lock()
	client.lastcall = "c"
	client.innerHandler["c"] = func(resp *Response) {
		defer mutex.Unlock()
		if resp.DataType == dtError {
			err = getError(resp.Data)
			return
		}
		handle = resp.Handle
	}
	id := IdGen.Id()
	req := getJob(id, []byte(funcname), data)
	req.DataType = flag
	if err = client.write(req); err != nil {
		delete(client.innerHandler, "c")
		client.lastcall = ""
		return
	}
	mutex.Lock()
	return
}

// Call the function and get a response.
// flag can be set to: JobLow, JobNormal and JobHigh
func (client *Client) Do(funcname string, data []byte,
	flag byte, h ResponseHandler) (handle string, err error) {
	var datatype uint32
	switch flag {
	case JobLow:
		datatype = dtSubmitJobLow
	case JobHigh:
		datatype = dtSubmitJobHigh
	default:
		datatype = dtSubmitJob
	}
	handle, err = client.do(funcname, data, datatype)
	if err == nil && h != nil {
		client.respHandler[handle] = h
	}
	return
}

// Call the function in background, no response needed.
// flag can be set to: JobLow, JobNormal and JobHigh
func (client *Client) DoBg(funcname string, data []byte,
	flag byte) (handle string, err error) {
	if client.conn == nil {
		return "", ErrLostConn
	}
	var datatype uint32
	switch flag {
	case JobLow:
		datatype = dtSubmitJobLowBg
	case JobHigh:
		datatype = dtSubmitJobHighBg
	default:
		datatype = dtSubmitJobBg
	}
	handle, err = client.do(funcname, data, datatype)
	return
}

// Get job status from job server.
func (client *Client) Status(handle string) (status *Status, err error) {
	if client.conn == nil {
		return nil, ErrLostConn
	}
	var mutex sync.Mutex
	mutex.Lock()
	client.lastcall = "s" + handle
	client.innerHandler["s"+handle] = func(resp *Response) {
		defer mutex.Unlock()
		var err error
		status, err = resp._status()
		if err != nil {
			client.err(err)
		}
	}
	req := getRequest()
	req.DataType = dtGetStatus
	req.Data = []byte(handle)
	client.write(req)
	mutex.Lock()
	return
}

// Echo.
func (client *Client) Echo(data []byte) (echo []byte, err error) {
	if client.conn == nil {
		return nil, ErrLostConn
	}
	var mutex sync.Mutex
	mutex.Lock()
	client.innerHandler["e"] = func(resp *Response) {
		echo = resp.Data
		mutex.Unlock()
	}
	req := getRequest()
	req.DataType = dtEchoReq
	req.Data = data
	client.lastcall = "e"
	client.write(req)
	mutex.Lock()
	return
}

// Close connection
func (client *Client) Close() (err error) {
	client.Lock()
	defer client.Unlock()
	if client.conn != nil {
		err = client.conn.Close()
		client.conn = nil
	}
	return
}
<commit_msg>Replace mutex in client.do() with a channel to avoid deadlock and introduce command timeout<commit_after>// The client package helps developers connect to Gearmand, send
// jobs and fetch result.
package client

import (
	"bufio"
	"net"
	"sync"
	"time"
)

// One client connect to one server.
// Use Pool for multi-connections.
type Client struct {
	sync.Mutex

	net, addr, lastcall string
	respHandler         map[string]ResponseHandler
	innerHandler        map[string]ResponseHandler
	in                  chan *Response
	conn                net.Conn
	rw                  *bufio.ReadWriter

	RespTimeout         time.Duration  // response timeout for do() in ms

	ErrorHandler ErrorHandler
}

// Return a client.
func New(network, addr string) (client *Client, err error) {
	client = &Client{
		net:          network,
		addr:         addr,
		respHandler:  make(map[string]ResponseHandler, queueSize),
		innerHandler: make(map[string]ResponseHandler, queueSize),
		in:           make(chan *Response, queueSize),
		RespTimeout:  1000,
	}
	client.conn, err = net.Dial(client.net, client.addr)
	if err != nil {
		return
	}
	client.rw = bufio.NewReadWriter(bufio.NewReader(client.conn),
		bufio.NewWriter(client.conn))
	go client.readLoop()
	go client.processLoop()
	return
}

func (client *Client) write(req *request) (err error) {
	var n int
	buf := req.Encode()
	for i := 0; i < len(buf); i += n {
		n, err = client.rw.Write(buf[i:])
		if err != nil {
			return
		}
	}
	return client.rw.Flush()
}

func (client *Client) read(length int) (data []byte, err error) {
	n := 0
	buf := getBuffer(bufferSize)
	// read until data can be unpacked
	for i := length; i > 0 || len(data) < minPacketLength; i -= n {
		if n, err = client.rw.Read(buf); err != nil {
			return
		}
		data = append(data, buf[0:n]...)
		if n < bufferSize {
			break
		}
	}
	return
}

func (client *Client) readLoop() {
	defer close(client.in)
	var data, leftdata []byte
	var err error
	var resp *Response
ReadLoop:
	for client.conn != nil {
		if data, err = client.read(bufferSize); err != nil {
			if opErr, ok := err.(*net.OpError); ok {
				if opErr.Timeout() {
					client.err(err)
				}
				if opErr.Temporary() {
					continue
				}
				break
			}
			client.err(err)
			// If it is unexpected error and the connection wasn't
			// closed by Gearmand, the client should close the conection
			// and reconnect to job server.
			client.Close()
			client.conn, err = net.Dial(client.net, client.addr)
			if err != nil {
				client.err(err)
				break
			}
			client.rw = bufio.NewReadWriter(bufio.NewReader(client.conn),
				bufio.NewWriter(client.conn))
			continue
		}
		if len(leftdata) > 0 { // some data left for processing
			data = append(leftdata, data...)
			leftdata = nil
		}
		for {
			l := len(data)
			if l < minPacketLength { // not enough data
				leftdata = data
				continue ReadLoop
			}
			if resp, l, err = decodeResponse(data); err != nil {
				leftdata = data[l:]
				continue ReadLoop
			} else {
				client.in <- resp
			}
			data = data[l:]
			if len(data) > 0 {
				continue
			}
			break
		}
	}
}

func (client *Client) processLoop() {
	for resp := range client.in {
		switch resp.DataType {
		case dtError:
			if client.lastcall != "" {
				resp = client.handleInner(client.lastcall, resp)
				client.lastcall = ""
			} else {
				client.err(getError(resp.Data))
			}
		case dtStatusRes:
			resp = client.handleInner("s"+resp.Handle, resp)
		case dtJobCreated:
			resp = client.handleInner("c", resp)
		case dtEchoRes:
			resp = client.handleInner("e", resp)
		case dtWorkData, dtWorkWarning, dtWorkStatus:
			resp = client.handleResponse(resp.Handle, resp)
		case dtWorkComplete, dtWorkFail, dtWorkException:
			client.handleResponse(resp.Handle, resp)
			delete(client.respHandler, resp.Handle)
		}
	}
}

func (client *Client) err(e error) {
	if client.ErrorHandler != nil {
		client.ErrorHandler(e)
	}
}

func (client *Client) handleResponse(key string, resp *Response) *Response {
	if h, ok := client.respHandler[key]; ok {
		h(resp)
		return nil
	}
	return resp
}

func (client *Client) handleInner(key string, resp *Response) *Response {
	if h, ok := client.innerHandler[key]; ok {
		h(resp)
		delete(client.innerHandler, key)
		return nil
	}
	return resp
}

type handleOrError struct {
	handle string
	err error
}

func (client *Client) do(funcname string, data []byte,
	flag uint32) (handle string, err error) {
	if client.conn == nil {
		return "", ErrLostConn
	}
	var result = make(chan handleOrError, 1)
	client.lastcall = "c"
	client.innerHandler["c"] = func(resp *Response) {
		if resp.DataType == dtError {
			err = getError(resp.Data)
			result <- handleOrError{"", err}
			return
		}
		handle = resp.Handle
		result <- handleOrError{handle, nil}
	}
	id := IdGen.Id()
	req := getJob(id, []byte(funcname), data)
	req.DataType = flag
	if err = client.write(req); err != nil {
		delete(client.innerHandler, "c")
		client.lastcall = ""
		return
	}
	var timer = time.After(client.RespTimeout * time.Millisecond)
	select {
	case ret := <-result:
		return ret.handle, ret.err
	case <-timer:
		delete(client.innerHandler, "c")
		client.lastcall = ""
		return "", ErrLostConn
	}
	return
}

// Call the function and get a response.
// flag can be set to: JobLow, JobNormal and JobHigh
func (client *Client) Do(funcname string, data []byte,
	flag byte, h ResponseHandler) (handle string, err error) {
	var datatype uint32
	switch flag {
	case JobLow:
		datatype = dtSubmitJobLow
	case JobHigh:
		datatype = dtSubmitJobHigh
	default:
		datatype = dtSubmitJob
	}
	handle, err = client.do(funcname, data, datatype)
	if err == nil && h != nil {
		client.respHandler[handle] = h
	}
	return
}

// Call the function in background, no response needed.
// flag can be set to: JobLow, JobNormal and JobHigh
func (client *Client) DoBg(funcname string, data []byte,
	flag byte) (handle string, err error) {
	if client.conn == nil {
		return "", ErrLostConn
	}
	var datatype uint32
	switch flag {
	case JobLow:
		datatype = dtSubmitJobLowBg
	case JobHigh:
		datatype = dtSubmitJobHighBg
	default:
		datatype = dtSubmitJobBg
	}
	handle, err = client.do(funcname, data, datatype)
	return
}

// Get job status from job server.
func (client *Client) Status(handle string) (status *Status, err error) {
	if client.conn == nil {
		return nil, ErrLostConn
	}
	var mutex sync.Mutex
	mutex.Lock()
	client.lastcall = "s" + handle
	client.innerHandler["s"+handle] = func(resp *Response) {
		defer mutex.Unlock()
		var err error
		status, err = resp._status()
		if err != nil {
			client.err(err)
		}
	}
	req := getRequest()
	req.DataType = dtGetStatus
	req.Data = []byte(handle)
	client.write(req)
	mutex.Lock()
	return
}

// Echo.
func (client *Client) Echo(data []byte) (echo []byte, err error) {
	if client.conn == nil {
		return nil, ErrLostConn
	}
	var mutex sync.Mutex
	mutex.Lock()
	client.innerHandler["e"] = func(resp *Response) {
		echo = resp.Data
		mutex.Unlock()
	}
	req := getRequest()
	req.DataType = dtEchoReq
	req.Data = data
	client.lastcall = "e"
	client.write(req)
	mutex.Lock()
	return
}

// Close connection
func (client *Client) Close() (err error) {
	client.Lock()
	defer client.Unlock()
	if client.conn != nil {
		err = client.conn.Close()
		client.conn = nil
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	kingpin "gopkg.in/alecthomas/kingpin.v2"

	"github.com/ShowMax/go-fqdn"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/promlog"
	"github.com/prometheus/common/promlog/flag"
	"github.com/robustperception/pushprox/util"
)

var (
	myFqdn      = kingpin.Flag("fqdn", "FQDN to register with").Default(fqdn.Get()).String()
	proxyURL    = kingpin.Flag("proxy-url", "Push proxy to talk to.").Required().String()
	caCertFile  = kingpin.Flag("tls.cacert", "<file> CA certificate to verify peer against").String()
	tlsCert     = kingpin.Flag("tls.cert", "<cert> Client certificate file").String()
	tlsKey      = kingpin.Flag("tls.key", "<key> Private key file").String()
	metricsAddr = kingpin.Flag("metrics-addr", "Serve Prometheus metrics at this address").Default(":9369").String()
)

var (
	scrapeErrorCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "pushprox_client_scrape_errors_total",
			Help: "Number of scrape errors",
		},
	)
	pushErrorCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "pushprox_client_push_errors_total",
			Help: "Number of push errors",
		},
	)
	pollErrorCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "pushprox_client_poll_errors_total",
			Help: "Number of poll errors",
		},
	)
)

func init() {
	prometheus.MustRegister(pushErrorCounter, pollErrorCounter, scrapeErrorCounter)
}

type Coordinator struct {
	logger log.Logger
}

func (c *Coordinator) doScrape(request *http.Request, client *http.Client) {
	logger := log.With(c.logger, "scrape_id", request.Header.Get("id"))
	timeout, _ := util.GetHeaderTimeout(request.Header)
	ctx, _ := context.WithTimeout(request.Context(), timeout)
	request = request.WithContext(ctx)
	// We cannot handle https requests at the proxy, as we would only
	// see a CONNECT, so use a URL parameter to trigger it.
	params := request.URL.Query()
	if params.Get("_scheme") == "https" {
		request.URL.Scheme = "https"
		params.Del("_scheme")
		request.URL.RawQuery = params.Encode()
	}

	scrapeResp, err := client.Do(request)
	if err != nil {
		msg := fmt.Sprintf("Failed to scrape %s: %s", request.URL.String(), err)
		level.Warn(logger).Log("msg", "Failed to scrape", "Request URL", request.URL.String(), "err", err)
		scrapeErrorCounter.Inc()
		resp := &http.Response{
			StatusCode: 500,
			Header:     http.Header{},
			Body:       ioutil.NopCloser(strings.NewReader(msg)),
		}
		err = c.doPush(resp, request, client)
		if err != nil {
			pushErrorCounter.Inc()
			level.Warn(logger).Log("msg", "Failed to push failed scrape response:", "err", err)
			return
		}
		level.Info(logger).Log("msg", "Pushed failed scrape response")
		return
	}
	level.Info(logger).Log("msg", "Retrieved scrape response")
	err = c.doPush(scrapeResp, request, client)
	if err != nil {
		pushErrorCounter.Inc()
		level.Warn(logger).Log("msg", "Failed to push scrape response:", "err", err)
		return
	}
	level.Info(logger).Log("msg", "Pushed scrape result")
}

// Report the result of the scrape back up to the proxy.
func (c *Coordinator) doPush(resp *http.Response, origRequest *http.Request, client *http.Client) error {
	resp.Header.Set("id", origRequest.Header.Get("id")) // Link the request and response
	// Remaining scrape deadline.
	deadline, _ := origRequest.Context().Deadline()
	resp.Header.Set("X-Prometheus-Scrape-Timeout", fmt.Sprintf("%f", float64(time.Until(deadline))/1e9))

	base, err := url.Parse(*proxyURL)
	if err != nil {
		return err
	}
	u, err := url.Parse("push")
	if err != nil {
		return err
	}
	url := base.ResolveReference(u)

	buf := &bytes.Buffer{}
	resp.Write(buf)
	request := &http.Request{
		Method:        "POST",
		URL:           url,
		Body:          ioutil.NopCloser(buf),
		ContentLength: int64(buf.Len()),
	}
	request = request.WithContext(origRequest.Context())
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	return nil
}

func loop(c Coordinator, t *http.Transport) error {
	client := &http.Client{Transport: t}
	base, err := url.Parse(*proxyURL)
	if err != nil {
		level.Error(c.logger).Log("msg", "Error parsing url:", "err", err)
		return errors.New("error parsing url")
	}
	u, err := url.Parse("poll")
	if err != nil {
		level.Error(c.logger).Log("msg", "Error parsing url:", "err", err)
		return errors.New("error parsing url poll")
	}
	url := base.ResolveReference(u)
	resp, err := client.Post(url.String(), "", strings.NewReader(*myFqdn))
	if err != nil {
		level.Error(c.logger).Log("msg", "Error polling:", "err", err)
		return errors.New("error polling")
	}
	defer resp.Body.Close()

	request, err := http.ReadRequest(bufio.NewReader(resp.Body))
	if err != nil {
		level.Error(c.logger).Log("msg", "Error reading request:", "err", err)
		return errors.New("error reading request")
	}
	level.Info(c.logger).Log("msg", "Got scrape request", "scrape_id", request.Header.Get("id"), "url", request.URL)

	request.RequestURI = ""

	go c.doScrape(request, client)

	return nil
}

func main() {
	allowedLevel := promlog.AllowedLevel{}
	flag.AddFlags(kingpin.CommandLine, &allowedLevel)
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()
	logger := promlog.New(allowedLevel)
	coordinator := Coordinator{logger: logger}
	if *proxyURL == "" {
		level.Error(coordinator.logger).Log("msg", "--proxy-url flag must be specified.")
		os.Exit(1)
	}
	// Make sure proxyURL ends with a single '/'
	*proxyURL = strings.TrimRight(*proxyURL, "/") + "/"
	level.Info(coordinator.logger).Log("msg", "URL and FQDN info", "proxy_url", *proxyURL, "fqdn", *myFqdn)

	tlsConfig := &tls.Config{}
	if *tlsCert != "" {
		cert, err := tls.LoadX509KeyPair(*tlsCert, *tlsKey)
		if err != nil {
			level.Error(coordinator.logger).Log("msg", "Certificate or Key is invalid", "err", err)
			os.Exit(1)
		}

		// Setup HTTPS client
		tlsConfig.Certificates = []tls.Certificate{cert}

		tlsConfig.BuildNameToCertificate()
	}

	if *caCertFile != "" {
		caCert, err := ioutil.ReadFile(*caCertFile)
		if err != nil {
			level.Error(coordinator.logger).Log("msg", "Not able to read cacert file", "err", err)
			os.Exit(1)
		}
		caCertPool := x509.NewCertPool()
		if ok := caCertPool.AppendCertsFromPEM(caCert); !ok {
			level.Error(coordinator.logger).Log("msg", "Failed to use cacert file as ca certificate")
			os.Exit(1)
		}

		tlsConfig.RootCAs = caCertPool
	}

	if *metricsAddr != "" {
		go func() {
			if err := http.ListenAndServe(*metricsAddr, promhttp.Handler()); err != nil {
				level.Warn(coordinator.logger).Log("msg", "ListenAndServe", "err", err)
			}
		}()
	}

	transport := &http.Transport{TLSClientConfig: tlsConfig}

	for {
		err := loop(coordinator, transport)
		if err != nil {
			pollErrorCounter.Inc()
			time.Sleep(time.Second) // Don't pound the server. TODO: Randomised exponential backoff.
		}
	}
}
<commit_msg>client: Use sane default values for the custom transport (#45)<commit_after>package main

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	kingpin "gopkg.in/alecthomas/kingpin.v2"

	"github.com/ShowMax/go-fqdn"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/promlog"
	"github.com/prometheus/common/promlog/flag"
	"github.com/robustperception/pushprox/util"
)

var (
	myFqdn      = kingpin.Flag("fqdn", "FQDN to register with").Default(fqdn.Get()).String()
	proxyURL    = kingpin.Flag("proxy-url", "Push proxy to talk to.").Required().String()
	caCertFile  = kingpin.Flag("tls.cacert", "<file> CA certificate to verify peer against").String()
	tlsCert     = kingpin.Flag("tls.cert", "<cert> Client certificate file").String()
	tlsKey      = kingpin.Flag("tls.key", "<key> Private key file").String()
	metricsAddr = kingpin.Flag("metrics-addr", "Serve Prometheus metrics at this address").Default(":9369").String()
)

var (
	scrapeErrorCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "pushprox_client_scrape_errors_total",
			Help: "Number of scrape errors",
		},
	)
	pushErrorCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "pushprox_client_push_errors_total",
			Help: "Number of push errors",
		},
	)
	pollErrorCounter = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "pushprox_client_poll_errors_total",
			Help: "Number of poll errors",
		},
	)
)

func init() {
	prometheus.MustRegister(pushErrorCounter, pollErrorCounter, scrapeErrorCounter)
}

type Coordinator struct {
	logger log.Logger
}

func (c *Coordinator) doScrape(request *http.Request, client *http.Client) {
	logger := log.With(c.logger, "scrape_id", request.Header.Get("id"))
	timeout, _ := util.GetHeaderTimeout(request.Header)
	ctx, _ := context.WithTimeout(request.Context(), timeout)
	request = request.WithContext(ctx)
	// We cannot handle https requests at the proxy, as we would only
	// see a CONNECT, so use a URL parameter to trigger it.
	params := request.URL.Query()
	if params.Get("_scheme") == "https" {
		request.URL.Scheme = "https"
		params.Del("_scheme")
		request.URL.RawQuery = params.Encode()
	}

	scrapeResp, err := client.Do(request)
	if err != nil {
		msg := fmt.Sprintf("Failed to scrape %s: %s", request.URL.String(), err)
		level.Warn(logger).Log("msg", "Failed to scrape", "Request URL", request.URL.String(), "err", err)
		scrapeErrorCounter.Inc()
		resp := &http.Response{
			StatusCode: 500,
			Header:     http.Header{},
			Body:       ioutil.NopCloser(strings.NewReader(msg)),
		}
		err = c.doPush(resp, request, client)
		if err != nil {
			pushErrorCounter.Inc()
			level.Warn(logger).Log("msg", "Failed to push failed scrape response:", "err", err)
			return
		}
		level.Info(logger).Log("msg", "Pushed failed scrape response")
		return
	}
	level.Info(logger).Log("msg", "Retrieved scrape response")
	err = c.doPush(scrapeResp, request, client)
	if err != nil {
		pushErrorCounter.Inc()
		level.Warn(logger).Log("msg", "Failed to push scrape response:", "err", err)
		return
	}
	level.Info(logger).Log("msg", "Pushed scrape result")
}

// Report the result of the scrape back up to the proxy.
func (c *Coordinator) doPush(resp *http.Response, origRequest *http.Request, client *http.Client) error {
	resp.Header.Set("id", origRequest.Header.Get("id")) // Link the request and response
	// Remaining scrape deadline.
	deadline, _ := origRequest.Context().Deadline()
	resp.Header.Set("X-Prometheus-Scrape-Timeout", fmt.Sprintf("%f", float64(time.Until(deadline))/1e9))

	base, err := url.Parse(*proxyURL)
	if err != nil {
		return err
	}
	u, err := url.Parse("push")
	if err != nil {
		return err
	}
	url := base.ResolveReference(u)

	buf := &bytes.Buffer{}
	resp.Write(buf)
	request := &http.Request{
		Method:        "POST",
		URL:           url,
		Body:          ioutil.NopCloser(buf),
		ContentLength: int64(buf.Len()),
	}
	request = request.WithContext(origRequest.Context())
	_, err = client.Do(request)
	if err != nil {
		return err
	}
	return nil
}

func loop(c Coordinator, t *http.Transport) error {
	client := &http.Client{Transport: t}
	base, err := url.Parse(*proxyURL)
	if err != nil {
		level.Error(c.logger).Log("msg", "Error parsing url:", "err", err)
		return errors.New("error parsing url")
	}
	u, err := url.Parse("poll")
	if err != nil {
		level.Error(c.logger).Log("msg", "Error parsing url:", "err", err)
		return errors.New("error parsing url poll")
	}
	url := base.ResolveReference(u)
	resp, err := client.Post(url.String(), "", strings.NewReader(*myFqdn))
	if err != nil {
		level.Error(c.logger).Log("msg", "Error polling:", "err", err)
		return errors.New("error polling")
	}
	defer resp.Body.Close()

	request, err := http.ReadRequest(bufio.NewReader(resp.Body))
	if err != nil {
		level.Error(c.logger).Log("msg", "Error reading request:", "err", err)
		return errors.New("error reading request")
	}
	level.Info(c.logger).Log("msg", "Got scrape request", "scrape_id", request.Header.Get("id"), "url", request.URL)

	request.RequestURI = ""

	go c.doScrape(request, client)

	return nil
}

func main() {
	allowedLevel := promlog.AllowedLevel{}
	flag.AddFlags(kingpin.CommandLine, &allowedLevel)
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()
	logger := promlog.New(allowedLevel)
	coordinator := Coordinator{logger: logger}
	if *proxyURL == "" {
		level.Error(coordinator.logger).Log("msg", "--proxy-url flag must be specified.")
		os.Exit(1)
	}
	// Make sure proxyURL ends with a single '/'
	*proxyURL = strings.TrimRight(*proxyURL, "/") + "/"
	level.Info(coordinator.logger).Log("msg", "URL and FQDN info", "proxy_url", *proxyURL, "fqdn", *myFqdn)

	tlsConfig := &tls.Config{}
	if *tlsCert != "" {
		cert, err := tls.LoadX509KeyPair(*tlsCert, *tlsKey)
		if err != nil {
			level.Error(coordinator.logger).Log("msg", "Certificate or Key is invalid", "err", err)
			os.Exit(1)
		}

		// Setup HTTPS client
		tlsConfig.Certificates = []tls.Certificate{cert}

		tlsConfig.BuildNameToCertificate()
	}

	if *caCertFile != "" {
		caCert, err := ioutil.ReadFile(*caCertFile)
		if err != nil {
			level.Error(coordinator.logger).Log("msg", "Not able to read cacert file", "err", err)
			os.Exit(1)
		}
		caCertPool := x509.NewCertPool()
		if ok := caCertPool.AppendCertsFromPEM(caCert); !ok {
			level.Error(coordinator.logger).Log("msg", "Failed to use cacert file as ca certificate")
			os.Exit(1)
		}

		tlsConfig.RootCAs = caCertPool
	}

	if *metricsAddr != "" {
		go func() {
			if err := http.ListenAndServe(*metricsAddr, promhttp.Handler()); err != nil {
				level.Warn(coordinator.logger).Log("msg", "ListenAndServe", "err", err)
			}
		}()
	}

	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
			DualStack: true,
		}).DialContext,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
		TLSClientConfig: tlsConfig,
	}

	for {
		err := loop(coordinator, transport)
		if err != nil {
			pollErrorCounter.Inc()
			time.Sleep(time.Second) // Don't pound the server. TODO: Randomised exponential backoff.
		}
	}
}
<|endoftext|>


<commit_before>/*
	Package client implements the class of a network client which can interact with a mix network.
*/

package client

import (
	"net"

	"anonymous-messaging/clientCore"
	"anonymous-messaging/networker"
	"anonymous-messaging/pki"
	"anonymous-messaging/config"
	"crypto/elliptic"
	"github.com/protobuf/proto"

	log "github.com/sirupsen/logrus"
	"fmt"
	"anonymous-messaging/helpers"
	"time"
	"math"
)

const (
	desiredRateParameter = 0.2
	pathLength           = 2
	ASSIGNE_FLAG = "\xA2"
	COMM_FLAG = "\xC6"
	TOKEN_FLAG = "xA9"
	PULL_FLAG = "\xFF"
	MAX_BUFFERQUEUE_SIZE = 10000
)

type ClientIt interface {
	networker.NetworkClient
	networker.NetworkServer
	SendMessage(message string, recipient config.MixConfig)
	ProcessPacket(packet []byte)
	Start()
	ReadInMixnetPKI()
	ReadInClientsPKI()
}

type Client struct {
	Host string
	Port string
	clientCore.CryptoClient

	Listener *net.TCPListener

	PkiDir string
	OtherClients []config.ClientConfig

	Config config.ClientConfig
	token []byte

	OutQueue chan []byte

}

/*
	Start function creates the loggers for capturing the info and error logs;
	it reads the network and users information from the PKI database
	and starts the listening server. Function returns an error
	signaling whenever any operation was unsuccessful.
*/
func (c *Client) Start() error {
	c.OutQueue = make(chan []byte)

	err := c.ReadInClientsPKI(c.PkiDir)
	if err != nil{
		return err
	}

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

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

	c.Run()
	return nil
}

/*
	SendMessage responsible for sending a real message. Takes as input the message string
	and the public information about the destination.
	The function generates a random path and a set of random values from exponential distribution.
	Given those values it triggers the encode function, which packs the message into the
	sphinx cryptographic packet format. Next, the encoded packet is combined with a
	flag signaling that this is a usual network packet, and passed to be send.
	The function returns an error if any issues occurred.
*/
func (c *Client) SendMessage(message string, recipient config.ClientConfig) error {

	sphinxPacket, err := c.CreateSphinxPacket(message, recipient)
	if err != nil {
		return err
	}

	packetBytes, err := config.WrapWithFlag(COMM_FLAG, sphinxPacket)
	if err != nil{
		return err
	}

	//err = c.Send(packetBytes, c.Provider.Host, c.Provider.Port)
	//if err != nil {
	//	return err
	//}
	c.OutQueue <- packetBytes
	return nil
}

/*
	Send opens a connection with selected network address
	and send the passed packet. If connection failed or
	the packet could not be send, an error is returned
*/
func (c *Client) Send(packet []byte, host string, port string) error {

	conn, err := net.Dial("tcp", host+":"+port)

	if err != nil {
		return err
	}
	defer conn.Close()

	_, err = conn.Write(packet)
	return err
}

/*
	ListenForIncomingConnections responsible for running the listening process of the server;
	The clients listener accepts incoming connections and
	passes the incoming packets to the packet handler.
	If the connection could not be accepted an error
	is logged into the log files, but the function is not stopped
*/
func (c *Client) ListenForIncomingConnections() {
	for {
		conn, err := c.Listener.Accept()

		if err != nil {
			log.WithFields(log.Fields{"id" : c.Id}).Error(err)
		} else {
			go c.HandleConnection(conn)
		}
	}
}

/*
	HandleConnection handles the received packets; it checks the flag of the
	packet and schedules a corresponding process function;
	The potential errors are logged into the log files.
*/
func (c *Client) HandleConnection(conn net.Conn) {

	buff := make([]byte, 1024)
	defer conn.Close()

	reqLen, err := conn.Read(buff)
	if err != nil {
		log.WithFields(log.Fields{"id" : c.Id}).Error(err)
		panic(err)
	}
	var packet config.GeneralPacket
	err = proto.Unmarshal(buff[:reqLen], &packet)
	if err != nil {
		log.WithFields(log.Fields{"id" : c.Id}).Error(err)
	}

	switch packet.Flag {
	case TOKEN_FLAG:
		c.RegisterToken(packet.Data)
		go func() {
			err = c.SendMessage("Hello world, this is me", c.OtherClients[0])
			if err != nil {
				log.WithFields(log.Fields{"id" : c.Id}).Error(err)
			}

			err = c.GetMessagesFromProvider()
			if err != nil {
				log.WithFields(log.Fields{"id" : c.Id}).Error(err)
			}
		}()
	case COMM_FLAG:
		_, err := c.ProcessPacket(packet.Data)
		if err != nil {
			log.WithFields(log.Fields{"id" : c.Id}).Error(err)
		}
		log.WithFields(log.Fields{"id" : c.Id}).Info(" Received new message")
	default:
		log.WithFields(log.Fields{"id" : c.Id}).Info(" Packet flag not recognised. Packet dropped.")
	}
}

/*
	RegisterToken stores the authentication token received from the provider
 */
func (c *Client) RegisterToken(token []byte) {
	c.token = token
	log.WithFields(log.Fields{"id" : c.Id}).Info(fmt.Sprintf(" Registered token %s", c.token))
}

/*
	ProcessPacket processes the received sphinx packet and returns the
	encapsulated message or error in case the processing
	was unsuccessful.
 */
func (c *Client) ProcessPacket(packet []byte) ([]byte, error) {
	log.WithFields(log.Fields{"id" : c.Id}).Info(" Processing packet")
	return packet, nil
}

/*
	RegisterToProvider allows the client to register with the selected provider.
	The client sends a special assignment packet, with its public information, to the provider
	or returns an error.
*/
func (c *Client) RegisterToProvider() error{

	log.WithFields(log.Fields{"id" : c.Id}).Info(" Sending request to provider to register")

	confBytes, err := proto.Marshal(&c.Config)
	if err != nil{
		return err
	}

	pktBytes, err := config.WrapWithFlag(ASSIGNE_FLAG, confBytes)
	if err != nil{
		return err
	}

	err = c.Send(pktBytes, c.Provider.Host, c.Provider.Port)
	if err != nil{
		return err
	}
	return nil
}

/*
	GetMessagesFromProvider allows to fetch messages from the inbox stored by the
	provider. The client sends a pull packet to the provider, along with
	the authentication token. An error is returned if occurred.
*/
func (c *Client) GetMessagesFromProvider() error {
	pullRqs := config.PullRequest{ClientId: c.Id, Token: c.token}
	pullRqsBytes, err := proto.Marshal(&pullRqs)
	if err != nil{
		return err
	}

	pktBytes, err := config.WrapWithFlag(PULL_FLAG, pullRqsBytes)
	if err != nil{
		return err
	}

	err = c.Send(pktBytes, c.Provider.Host, c.Provider.Port)
	if err != nil{
		return err
	}

	return nil
}

/*
	Run opens the listener to start listening on clients host and port
 */
func (c *Client) Run() {
	defer c.Listener.Close()
	finish := make(chan bool)

	go func() {
		c.ControlOutQueue()
	}()

	go func() {
		c.FakeAdding()
	}()

	go func() {
		log.WithFields(log.Fields{"id" : c.Id}).Info(fmt.Sprintf("Listening on address %s", c.Host + ":" + c.Port))
		c.ListenForIncomingConnections()
	}()
	<-finish
}

func (c *Client) FakeAdding(){
	log.WithFields(log.Fields{"id" : c.Id}).Info("Started fake adding")
	for {
		packet := []byte("Hello world")
		c.OutQueue <- packet
		log.WithFields(log.Fields{"id" : c.Id}).Info("Added packet")
		time.Sleep(10 * time.Second)
	}
}

func (c *Client) ControlOutQueue() error{
	log.WithFields(log.Fields{"id" : c.Id}).Info("Queue controller started")
	for{
		select{
		case realPacket := <-c.OutQueue:
			c.Send(realPacket, c.Provider.Host, c.Provider.Port)
			log.WithFields(log.Fields{"id" : c.Id}).Info("Real packet was send")
			delaySec, err := helpers.RandomExponential(desiredRateParameter)
			if err != nil{
				return err
			}
			time.Sleep(time.Duration(int64(delaySec * math.Pow10(9))) * time.Nanosecond)
		default:
			delaySec, err := helpers.RandomExponential(desiredRateParameter)
			if err != nil{
				return err
			}
			log.WithFields(log.Fields{"id" : c.Id}).Info("OutQueue empty. Dummy packet sent.")
			time.Sleep(time.Duration(int64(delaySec * math.Pow10(9))) * time.Nanosecond)
		}
	}
	return nil
}

/*
	ReadInMixnetPKI reads in the public information about active mixes
	from the PKI database and stores them locally. In case
	the connection or fetching data from the PKI went wrong,
	an error is returned.
*/
func (c *Client) ReadInMixnetPKI(pkiName string) error {
	log.WithFields(log.Fields{"id" : c.Id}).Info(fmt.Sprintf("Reading network information from the PKI: %s", pkiName))

	db, err := pki.OpenDatabase(pkiName, "sqlite3")

	if err != nil{
		return err
	}

	records, err := pki.QueryDatabase(db, "Pki", "Mix")

	if err != nil{
		return err
	}

	for records.Next() {
		result := make(map[string]interface{})
		err := records.MapScan(result)

		if err != nil {
			return err
		}

		var pubs config.MixConfig
		err = proto.Unmarshal(result["Config"].([]byte), &pubs)
		if err != nil {
			return err
		}

		c.ActiveMixes = append(c.ActiveMixes, pubs)
	}

	log.WithFields(log.Fields{"id" : c.Id}).Info(" Network information uploaded")
	return nil
}

/*
	ReadInClientsPKI reads in the public information about users
	from the PKI database and stores them locally. In case
	the connection or fetching data from the PKI went wrong,
	an error is returned.
*/
func (c *Client) ReadInClientsPKI(pkiName string) error {
	log.WithFields(log.Fields{"id" : c.Id}).Info(fmt.Sprintf(" Reading network users information from the PKI: %s", pkiName))

	db, err := pki.OpenDatabase(pkiName, "sqlite3")

	if err != nil{
		return err
	}

	records, err := pki.QueryDatabase(db, "Pki", "Client")

	if err != nil {
		return err
	}

	for records.Next() {
		result := make(map[string]interface{})
		err := records.MapScan(result)

		if err != nil {
			return err
		}

		var pubs config.ClientConfig
		err = proto.Unmarshal(result["Config"].([]byte), &pubs)
		if err != nil {
			return err
		}
		c.OtherClients = append(c.OtherClients, pubs)
	}
	log.WithFields(log.Fields{"id" : c.Id}).Info("  Information about other users uploaded")
	return nil
}

/*
	The constructor function to create an new client object.
	Function returns a new client object or an error, if occurred.
*/
func NewClient(id, host, port string, pubKey []byte, prvKey []byte, pkiDir string, provider config.MixConfig) (*Client, error) {
	core := clientCore.CryptoClient{Id: id, PubKey: pubKey, PrvKey: prvKey, Curve: elliptic.P224(), Provider: provider}

	c := Client{Host: host, Port: port, CryptoClient: core, PkiDir: pkiDir}
	c.Config = config.ClientConfig{Id : c.Id, Host: c.Host, Port: c.Port, PubKey: c.PubKey, Provider: &c.Provider}

	configBytes, err := proto.Marshal(&c.Config)

	if err != nil{
		return nil, err
	}
	err = helpers.AddToDatabase(pkiDir, "Pki", c.Id, "Client", configBytes)
	if err != nil{
		return nil, err
	}

	addr, err := helpers.ResolveTCPAddress(c.Host, c.Port)
	if err != nil {
		return nil, err
	}

	c.Listener, err = net.ListenTCP("tcp", addr)
	if err != nil {
		return nil, err
	}
	return &c, nil
}

/*
	NewTestClient constructs a client object, which can be used for testing. The object contains the crypto core
	and the top-level of client, but does not involve networking and starting a listener.
 */
func NewTestClient(id, host, port string, pubKey []byte, prvKey []byte, pkiDir string, provider config.MixConfig) (*Client, error) {
	core := clientCore.CryptoClient{Id: id, PubKey: pubKey, PrvKey: prvKey, Curve: elliptic.P224(), Provider: provider}
	c := Client{Host: host, Port: port, CryptoClient: core, PkiDir: pkiDir}
	c.Config = config.ClientConfig{Id : c.Id, Host: c.Host, Port: c.Port, PubKey: c.PubKey, Provider: &c.Provider}

	return &c, nil
}<commit_msg>Added loop cover message for empty queue<commit_after>/*
	Package client implements the class of a network client which can interact with a mix network.
*/

package client

import (
	"net"

	"anonymous-messaging/clientCore"
	"anonymous-messaging/networker"
	"anonymous-messaging/pki"
	"anonymous-messaging/config"
	"crypto/elliptic"
	"github.com/protobuf/proto"

	log "github.com/sirupsen/logrus"
	"fmt"
	"anonymous-messaging/helpers"
	"time"
	"math"
)

const (
	desiredRateParameter = 0.2
	pathLength           = 2
	ASSIGNE_FLAG = "\xA2"
	COMM_FLAG = "\xC6"
	TOKEN_FLAG = "xA9"
	PULL_FLAG = "\xFF"
	MAX_BUFFERQUEUE_SIZE = 10000
)

type ClientIt interface {
	networker.NetworkClient
	networker.NetworkServer
	SendMessage(message string, recipient config.MixConfig)
	ProcessPacket(packet []byte)
	Start()
	ReadInMixnetPKI()
	ReadInClientsPKI()
}

type Client struct {
	Host string
	Port string
	clientCore.CryptoClient

	Listener *net.TCPListener

	PkiDir string
	OtherClients []config.ClientConfig

	Config config.ClientConfig
	token []byte

	OutQueue chan []byte

}

/*
	Start function creates the loggers for capturing the info and error logs;
	it reads the network and users information from the PKI database
	and starts the listening server. Function returns an error
	signaling whenever any operation was unsuccessful.
*/
func (c *Client) Start() error {
	c.OutQueue = make(chan []byte)

	err := c.ReadInClientsPKI(c.PkiDir)
	if err != nil{
		return err
	}

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

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

	c.Run()
	return nil
}

/*
	SendMessage responsible for sending a real message. Takes as input the message string
	and the public information about the destination.
	The function generates a random path and a set of random values from exponential distribution.
	Given those values it triggers the encode function, which packs the message into the
	sphinx cryptographic packet format. Next, the encoded packet is combined with a
	flag signaling that this is a usual network packet, and passed to be send.
	The function returns an error if any issues occurred.
	TODO change message type to []byte
*/
func (c *Client) SendMessage(message string, recipient config.ClientConfig) error {

	sphinxPacket, err := c.CreateSphinxPacket(message, recipient)
	if err != nil {
		return err
	}

	packetBytes, err := config.WrapWithFlag(COMM_FLAG, sphinxPacket)
	if err != nil{
		return err
	}

	//err = c.Send(packetBytes, c.Provider.Host, c.Provider.Port)
	//if err != nil {
	//	return err
	//}
	c.OutQueue <- packetBytes
	return nil
}

/*
	Send opens a connection with selected network address
	and send the passed packet. If connection failed or
	the packet could not be send, an error is returned
*/
func (c *Client) Send(packet []byte, host string, port string) error {

	conn, err := net.Dial("tcp", host+":"+port)

	if err != nil {
		return err
	}
	defer conn.Close()

	_, err = conn.Write(packet)
	return err
}

/*
	ListenForIncomingConnections responsible for running the listening process of the server;
	The clients listener accepts incoming connections and
	passes the incoming packets to the packet handler.
	If the connection could not be accepted an error
	is logged into the log files, but the function is not stopped
*/
func (c *Client) ListenForIncomingConnections() {
	for {
		conn, err := c.Listener.Accept()

		if err != nil {
			log.WithFields(log.Fields{"id" : c.Id}).Error(err)
		} else {
			go c.HandleConnection(conn)
		}
	}
}

/*
	HandleConnection handles the received packets; it checks the flag of the
	packet and schedules a corresponding process function;
	The potential errors are logged into the log files.
*/
func (c *Client) HandleConnection(conn net.Conn) {

	buff := make([]byte, 1024)
	defer conn.Close()

	reqLen, err := conn.Read(buff)
	if err != nil {
		log.WithFields(log.Fields{"id" : c.Id}).Error(err)
		panic(err)
	}
	var packet config.GeneralPacket
	err = proto.Unmarshal(buff[:reqLen], &packet)
	if err != nil {
		log.WithFields(log.Fields{"id" : c.Id}).Error(err)
	}

	switch packet.Flag {
	case TOKEN_FLAG:
		c.RegisterToken(packet.Data)
		go func() {
			err = c.SendMessage("Hello world, this is me", c.OtherClients[0])
			if err != nil {
				log.WithFields(log.Fields{"id" : c.Id}).Error(err)
			}

			err = c.GetMessagesFromProvider()
			if err != nil {
				log.WithFields(log.Fields{"id" : c.Id}).Error(err)
			}
		}()
	case COMM_FLAG:
		_, err := c.ProcessPacket(packet.Data)
		if err != nil {
			log.WithFields(log.Fields{"id" : c.Id}).Error(err)
		}
		log.WithFields(log.Fields{"id" : c.Id}).Info(" Received new message")
	default:
		log.WithFields(log.Fields{"id" : c.Id}).Info(" Packet flag not recognised. Packet dropped.")
	}
}

/*
	RegisterToken stores the authentication token received from the provider
 */
func (c *Client) RegisterToken(token []byte) {
	c.token = token
	log.WithFields(log.Fields{"id" : c.Id}).Info(fmt.Sprintf(" Registered token %s", c.token))
}

/*
	ProcessPacket processes the received sphinx packet and returns the
	encapsulated message or error in case the processing
	was unsuccessful.
 */
func (c *Client) ProcessPacket(packet []byte) ([]byte, error) {
	log.WithFields(log.Fields{"id" : c.Id}).Info(" Processing packet")
	return packet, nil
}

/*
	RegisterToProvider allows the client to register with the selected provider.
	The client sends a special assignment packet, with its public information, to the provider
	or returns an error.
*/
func (c *Client) RegisterToProvider() error{

	log.WithFields(log.Fields{"id" : c.Id}).Info(" Sending request to provider to register")

	confBytes, err := proto.Marshal(&c.Config)
	if err != nil{
		return err
	}

	pktBytes, err := config.WrapWithFlag(ASSIGNE_FLAG, confBytes)
	if err != nil{
		return err
	}

	err = c.Send(pktBytes, c.Provider.Host, c.Provider.Port)
	if err != nil{
		return err
	}
	return nil
}

/*
	GetMessagesFromProvider allows to fetch messages from the inbox stored by the
	provider. The client sends a pull packet to the provider, along with
	the authentication token. An error is returned if occurred.
*/
func (c *Client) GetMessagesFromProvider() error {
	pullRqs := config.PullRequest{ClientId: c.Id, Token: c.token}
	pullRqsBytes, err := proto.Marshal(&pullRqs)
	if err != nil{
		return err
	}

	pktBytes, err := config.WrapWithFlag(PULL_FLAG, pullRqsBytes)
	if err != nil{
		return err
	}

	err = c.Send(pktBytes, c.Provider.Host, c.Provider.Port)
	if err != nil{
		return err
	}

	return nil
}

/*
	Run opens the listener to start listening on clients host and port
 */
func (c *Client) Run() {
	defer c.Listener.Close()
	finish := make(chan bool)

	go func() {
		c.ControlOutQueue()
	}()

	go func() {
		c.FakeAdding()
	}()

	go func() {
		log.WithFields(log.Fields{"id" : c.Id}).Info(fmt.Sprintf("Listening on address %s", c.Host + ":" + c.Port))
		c.ListenForIncomingConnections()
	}()
	<-finish
}

func (c *Client) FakeAdding(){
	log.WithFields(log.Fields{"id" : c.Id}).Info("Started fake adding")
	for {
		packet := []byte("Hello world")
		c.OutQueue <- packet
		log.WithFields(log.Fields{"id" : c.Id}).Info("Added packet")
		time.Sleep(10 * time.Second)
	}
}

func (c *Client) ControlOutQueue() error{
	log.WithFields(log.Fields{"id" : c.Id}).Info("Queue controller started")
	for{
		select{
		case realPacket := <-c.OutQueue:
			c.Send(realPacket, c.Provider.Host, c.Provider.Port)
			log.WithFields(log.Fields{"id" : c.Id}).Info("Real packet was send")
			delaySec, err := helpers.RandomExponential(desiredRateParameter)
			if err != nil{
				return err
			}
			time.Sleep(time.Duration(int64(delaySec * math.Pow10(9))) * time.Nanosecond)
		default:
			delaySec, err := helpers.RandomExponential(desiredRateParameter)
			if err != nil{
				return err
			}
			dummyPacket, err := c.CreateCoverMessage()
			if err != nil{
				return err
			}
			c.Send(dummyPacket, c.Provider.Host, c.Provider.Port)
			log.WithFields(log.Fields{"id" : c.Id}).Info("OutQueue empty. Dummy packet sent.")
			time.Sleep(time.Duration(int64(delaySec * math.Pow10(9))) * time.Nanosecond)
		}
	}
	return nil
}

/*
	CreateCoverMessage packs a dummy message into a Sphinx packet.
	The dummy message is a loop message.
	TODO: change to a drop cover message.
 */
func (c *Client) CreateCoverMessage() ([]byte, error) {
	dummyLoad := "DummyPayloadMessage"
	sphinxPacket, err := c.CreateSphinxPacket(dummyLoad, c.Config)
	if err != nil {
		return nil, err
	}

	packetBytes, err := config.WrapWithFlag(COMM_FLAG, sphinxPacket)
	if err != nil{
		return nil, err
	}
	return packetBytes, nil
}

/*
	ReadInMixnetPKI reads in the public information about active mixes
	from the PKI database and stores them locally. In case
	the connection or fetching data from the PKI went wrong,
	an error is returned.
*/
func (c *Client) ReadInMixnetPKI(pkiName string) error {
	log.WithFields(log.Fields{"id" : c.Id}).Info(fmt.Sprintf("Reading network information from the PKI: %s", pkiName))

	db, err := pki.OpenDatabase(pkiName, "sqlite3")

	if err != nil{
		return err
	}

	records, err := pki.QueryDatabase(db, "Pki", "Mix")

	if err != nil{
		return err
	}

	for records.Next() {
		result := make(map[string]interface{})
		err := records.MapScan(result)

		if err != nil {
			return err
		}

		var pubs config.MixConfig
		err = proto.Unmarshal(result["Config"].([]byte), &pubs)
		if err != nil {
			return err
		}

		c.ActiveMixes = append(c.ActiveMixes, pubs)
	}

	log.WithFields(log.Fields{"id" : c.Id}).Info(" Network information uploaded")
	return nil
}

/*
	ReadInClientsPKI reads in the public information about users
	from the PKI database and stores them locally. In case
	the connection or fetching data from the PKI went wrong,
	an error is returned.
*/
func (c *Client) ReadInClientsPKI(pkiName string) error {
	log.WithFields(log.Fields{"id" : c.Id}).Info(fmt.Sprintf(" Reading network users information from the PKI: %s", pkiName))

	db, err := pki.OpenDatabase(pkiName, "sqlite3")

	if err != nil{
		return err
	}

	records, err := pki.QueryDatabase(db, "Pki", "Client")

	if err != nil {
		return err
	}

	for records.Next() {
		result := make(map[string]interface{})
		err := records.MapScan(result)

		if err != nil {
			return err
		}

		var pubs config.ClientConfig
		err = proto.Unmarshal(result["Config"].([]byte), &pubs)
		if err != nil {
			return err
		}
		c.OtherClients = append(c.OtherClients, pubs)
	}
	log.WithFields(log.Fields{"id" : c.Id}).Info("  Information about other users uploaded")
	return nil
}

/*
	The constructor function to create an new client object.
	Function returns a new client object or an error, if occurred.
*/
func NewClient(id, host, port string, pubKey []byte, prvKey []byte, pkiDir string, provider config.MixConfig) (*Client, error) {
	core := clientCore.CryptoClient{Id: id, PubKey: pubKey, PrvKey: prvKey, Curve: elliptic.P224(), Provider: provider}

	c := Client{Host: host, Port: port, CryptoClient: core, PkiDir: pkiDir}
	c.Config = config.ClientConfig{Id : c.Id, Host: c.Host, Port: c.Port, PubKey: c.PubKey, Provider: &c.Provider}

	configBytes, err := proto.Marshal(&c.Config)

	if err != nil{
		return nil, err
	}
	err = helpers.AddToDatabase(pkiDir, "Pki", c.Id, "Client", configBytes)
	if err != nil{
		return nil, err
	}

	addr, err := helpers.ResolveTCPAddress(c.Host, c.Port)
	if err != nil {
		return nil, err
	}

	c.Listener, err = net.ListenTCP("tcp", addr)
	if err != nil {
		return nil, err
	}
	return &c, nil
}

/*
	NewTestClient constructs a client object, which can be used for testing. The object contains the crypto core
	and the top-level of client, but does not involve networking and starting a listener.
 */
func NewTestClient(id, host, port string, pubKey []byte, prvKey []byte, pkiDir string, provider config.MixConfig) (*Client, error) {
	core := clientCore.CryptoClient{Id: id, PubKey: pubKey, PrvKey: prvKey, Curve: elliptic.P224(), Provider: provider}
	c := Client{Host: host, Port: port, CryptoClient: core, PkiDir: pkiDir}
	c.Config = config.ClientConfig{Id : c.Id, Host: c.Host, Port: c.Port, PubKey: c.PubKey, Provider: &c.Provider}

	return &c, nil
}<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"fmt"
	"net/url"
	"os"
	"sync"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	"github.com/gorilla/websocket"
	"github.com/mohae/autofact"
	"github.com/mohae/autofact/message"
	"github.com/mohae/autofact/sysinfo"
)

// Client is anything that talks to the server.
type Client struct {
	// Inf holds the basic client information.
	*Inf
	// This is Inf as bytes: this is hopefully unnecessary, but I don't know at this point.
	InfBytes []byte
	// Conn holds the configuration for connecting to the server.
	ConnCfg `json:"-"`
	// Cfg holds the client configuration (how the client behaves).
	Cfg `json:"-"`

	// Healthbeat buffers
	CPUData [][]byte `json:"cpu_stat"`
	MemData [][]byte `json:"mem_data"`

	muSend sync.Mutex
	WS     *websocket.Conn `json:"-"`
	// Channel for outbound binary messages.  The message is assumed to be a
	// websocket.Binary type
	SendB       chan []byte `json:"-"`
	SendStr     chan string `json:"-"`
	mu          sync.Mutex
	isConnected bool
	ServerURL   url.URL `json:"-"`
}

func New(id uint32, name string) *Client {
	bldr := flatbuffers.NewBuilder(0)
	n := bldr.CreateString(name)
	InfStart(bldr)
	InfAddID(bldr, id)
	InfAddHostname(bldr, n)
	bldr.Finish(InfEnd(bldr))
	return &Client{
		Inf:      GetRootAsInf(bldr.Bytes[bldr.Head():], 0),
		InfBytes: bldr.Bytes[bldr.Head():],
		// A really small buffer:
		// TODO: rethink this vis-a-vis what happens when recipient isn't there
		// or if it goes away during sending and possibly caching items to be sent.
		SendB:   make(chan []byte, 10),
		SendStr: make(chan string, 10),
	}
}

// Connect handles connecting to the server and returns the connection status.
// The client will attempt to connect until it has either succeeded or the
// connection retry period has been exceeded.  A retry is done every 5 seconds.
//
// If the client is already connected, nothing will be done.
func (c *Client) Connect() bool {
	// If already connected, return that fact.
	if c.IsConnected() {
		return true
	}
	start := time.Now()
	retryEnd := start.Add(c.ConnectPeriod)
	// connect to server; retry until the retry period has expired
	for {
		if time.Now().After(retryEnd) {
			fmt.Fprintf(os.Stderr, "timed out while trying to connect to the server: %s\n", c.ServerURL.String())
			return false
		}
		err := c.DialServer()
		if err == nil {
			break
		}
		time.Sleep(c.ConnectInterval)
		fmt.Printf("unable to connect to the server %s: retrying...\n", c.ServerURL.String())
	}
	// Send the ClientInf.  If the ID == 0 or it can't be found, the server will
	// respond with one.  Retry until the server responds, or until the
	// reconnectPeriod has expired.
	err := c.WS.WriteMessage(websocket.BinaryMessage, c.InfBytes)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while sending ID: %s\n", err)
		c.WS.Close()
		return false
	}

	typ, p, err := c.WS.ReadMessage()
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while Reading ID response: %s\n", err)
		c.WS.Close()
		return false
	}

	switch typ {
	case websocket.BinaryMessage:
		// a binary response is ClientInf
		c.Inf = GetRootAsInf(p, 0)
		c.InfBytes = p
		fmt.Printf("new ID: %d\n", c.Inf.ID())
	case websocket.TextMessage:
		fmt.Printf("%s\n", string(p))
	default:
		fmt.Printf("unexpected welcome response type %d: %v\n", typ, p)
		c.WS.Close()
		return false
	}
	// the next message is the current Cfg, which is applied to the clientID
	typ, p, err = c.WS.ReadMessage()
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while Reading ID response: %s\n", err)
		c.WS.Close()
		return false
	}

	switch typ {
	case websocket.BinaryMessage:
		// decode the message (it's in flatbuffer format), and apply to the cfg
		c.processBinaryMessage(p)
	default:
		fmt.Printf("unexpected message type from the server: was expecting the client's cfg: %d: %v\n", typ, p)
		fmt.Println("closing the connection and exiting")
		c.WS.Close()
		return false
	}

	c.mu.Lock()
	c.isConnected = true
	c.mu.Unlock()
	return true
}

func (c *Client) DialServer() error {
	var err error
	c.WS, _, err = websocket.DefaultDialer.Dial(c.ServerURL.String(), nil)
	return err
}

func (c *Client) MessageWriter(doneCh chan struct{}) {
	pingPeriod := time.Duration(c.Cfg.PingPeriod())
	defer close(doneCh)
	for {
		select {
		case p, ok := <-c.SendB:
			// don't send if not connected
			if !c.IsConnected() {
				continue
			}
			if !ok {
				c.WS.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			err := c.WS.WriteMessage(websocket.BinaryMessage, p)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error writing message: %s\n", err)
			}
		case <-time.After(pingPeriod):
			// only ping if we are connected
			if !c.IsConnected() {
				continue
			}
			err := c.WS.WriteMessage(websocket.PingMessage, []byte("ping"))
			if err != nil {
				fmt.Fprintf(os.Stderr, "ping error: %s\n", err)
			}
		}
	}
}

func (c *Client) Reconnect() bool {
	c.mu.Lock()
	c.isConnected = false
	c.mu.Unlock()
	for i := 0; i < 4; i++ {
		b := c.Connect()
		if b {
			fmt.Println("reconnect true")
			return b
		}
	}
	return false
}

func (c *Client) Listen(doneCh chan struct{}) {
	// loop until there's a done signal
	defer close(doneCh)
	for {
		typ, p, err := c.WS.ReadMessage()
		if err != nil {
			fmt.Fprintf(os.Stderr, "error reading message: %s\n", err)
			if _, ok := err.(*websocket.CloseError); !ok {
				return
			}
			fmt.Println("reconnecting from read messages")
			connected := c.Reconnect()
			if connected {
				continue
			}
			fmt.Fprint(os.Stderr, "unable to reconnect to server")
			return
		}
		switch typ {
		case websocket.TextMessage:
			fmt.Printf("textmessage: %s\n", p)
			if bytes.Equal(p, autofact.AckMsg) {
				// if this is an acknowledgement message, do nothing
				continue
			}
			err := c.WS.WriteMessage(websocket.TextMessage, autofact.AckMsg)
			if err != nil {
				if _, ok := err.(*websocket.CloseError); !ok {
					return
				}
				fmt.Println("reconnect from writing message: textmessage")
				connected := c.Reconnect()
				if connected {
					continue
				}
				fmt.Fprint(os.Stderr, "unable to reconnect to server")
				return
			}
		case websocket.BinaryMessage:
			err = c.WS.WriteMessage(websocket.TextMessage, autofact.AckMsg)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error writing binary message: %s\n", err)
				if _, ok := err.(*websocket.CloseError); !ok {
					return
				}
				fmt.Println("reconnect from writing message: binarymessage")
				connected := c.Reconnect()
				if connected {
					continue
				}
				fmt.Fprint(os.Stderr, "unable to reconnect to server")
				return
			}
			c.processBinaryMessage(p)
		case websocket.CloseMessage:
			fmt.Printf("closemessage: %x\n", p)
			fmt.Println("reconnect from writing message: closemessage")
			connected := c.Reconnect()
			if connected {
				continue
			}
			fmt.Fprint(os.Stderr, "unable to reconnect to server")
			return
		}
	}
}

// IsConnected returns if the client is connected.
func (c *Client) IsConnected() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.isConnected
}

func (c *Client) PingHandler(msg string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	fmt.Printf("ping: %s\n", msg)
	return c.WS.WriteMessage(websocket.PongMessage, []byte("ping"))
}

func (c *Client) PongHandler(msg string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	fmt.Printf("pong: %s\n", msg)
	return c.WS.WriteMessage(websocket.PingMessage, []byte("pong"))
}

// TODO: should CPUData be enclosed in a struct for locking purposes?
func (c *Client) EnqueueCPUData(data []byte) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.CPUData = append(c.CPUData, data)
	return len(c.CPUData)
}

// EnqueueMemData adds the received data to the MemData buffer.  The current
// number of entries in the buffer is returned.
func (c *Client) EnqueueMemData(data []byte) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.MemData = append(c.MemData, data)
	return len(c.MemData)
}

// If the message send fails, whatever was cached will be lost.
// TODO: the current stats get copied for send; should the stat slice
// get reset now?  There is possible data loss this way; but there's
// possible data loss if the stat slice gets appended to between this
// copy op and the send completing, unless a lock is held the entire time,
// which could block the stats reading process leading to data loss due to
// missed reads.  I'm thinking COW or delete now; but punting becasue it
// really doesn't matter as this is just an experiment, right now.
// This also applies to CPUDatasFB
// TODO: should the messages to be sent be copied to a send cache so that
// there isn't data loss on a failed send?  Consecutive PushPeriods that
// failed to send may be problematic in that situation.
func (c *Client) FlushCPUData() [][]byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	data := make([][]byte, len(c.CPUData))
	copy(data, c.CPUData)
	c.CPUData = nil
	return data
}

// FlushMemData returns a copy of the client's MemData and nils the client's
// cache.
// The notes above apply here too.
func (c *Client) FlushMemData() [][]byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	data := make([][]byte, len(c.MemData))
	copy(data, c.MemData)
	c.MemData = nil
	return data
}

func (c *Client) Healthbeat() {
	// An interval of 0 means no healthbeat
	if c.Cfg.HealthbeatInterval() == 0 {
		return
	}
	cpuCh := make(chan []byte)
	memCh := make(chan []byte)
	go sysinfo.CPUDataTicker(time.Duration(c.Cfg.HealthbeatInterval()), cpuCh)
	go sysinfo.MemDataTicker(time.Duration(c.Cfg.HealthbeatInterval()), memCh)
	t := time.NewTicker(time.Duration(c.Cfg.HealthbeatPushPeriod()))
	defer t.Stop()
	for {
		select {
		case data, ok := <-cpuCh:
			if !ok {
				fmt.Println("cpu stats chan closed")
				goto done
			}
			c.EnqueueCPUData(data)
		case data, ok := <-memCh:
			if !ok {
				fmt.Println("cpu stats chan closed")
				goto done
			}
			c.EnqueueMemData(data)
		case <-t.C:
			if !c.IsConnected() {
				continue
			}
			c.SendData(message.CPUData, c.FlushCPUData())
			c.SendData(message.MemData, c.FlushMemData())
		}
	}
done:
	// Flush the buffer.
	c.SendData(message.CPUData, c.FlushCPUData())
	c.SendData(message.MemData, c.FlushMemData())
}

// SendData sends the received data to the server.  The caller checks to see
// if the client is connected to the server before calling.  If the connection
// is lost during processing, the cached stats will be lost.
// TODO:  should this be more resilient?
func (c *Client) SendData(kind message.Kind, data [][]byte) error {
	// for each data, send a message
	for _, v := range data {
		c.SendB <- message.Serialize(c.Inf.ID(), kind, v)
	}
	return nil
}

// SendMessage sends a single serialized message of type Kind.
func (c *Client) SendMessage(kind message.Kind, p []byte) {
	c.SendB <- message.Serialize(c.Inf.ID(), kind, p)
}

// binary messages are expected to be flatbuffer encoding of message.Message.
func (c *Client) processBinaryMessage(p []byte) error {
	// unmarshal the message
	msg := message.GetRootAsMessage(p, 0)
	// process according to kind
	k := message.Kind(msg.Kind())
	switch k {
	case message.ClientCfg:
		c.Cfg.Deserialize(msg.DataBytes())
	default:
		fmt.Println("unknown message kind")
		fmt.Println(string(p))
	}
	return nil
}

// Serialize serializes the Inf using flatbuffers and returns the []byte.
func (i *Inf) Serialize() []byte {
	bldr := flatbuffers.NewBuilder(0)
	h := bldr.CreateByteString(i.Hostname())
	r := bldr.CreateByteString(i.Region())
	z := bldr.CreateByteString(i.Zone())
	d := bldr.CreateByteString(i.DC())
	InfStart(bldr)
	InfAddID(bldr, i.ID())
	InfAddHostname(bldr, h)
	InfAddRegion(bldr, r)
	InfAddZone(bldr, z)
	InfAddDC(bldr, d)
	bldr.Finish(InfEnd(bldr))
	return bldr.Bytes[bldr.Head():]
}
<commit_msg>remove json tags from Client<commit_after>package client

import (
	"bytes"
	"fmt"
	"net/url"
	"os"
	"sync"
	"time"

	flatbuffers "github.com/google/flatbuffers/go"
	"github.com/gorilla/websocket"
	"github.com/mohae/autofact"
	"github.com/mohae/autofact/message"
	"github.com/mohae/autofact/sysinfo"
)

// Client is anything that talks to the server.
type Client struct {
	// Inf holds the basic client information.
	*Inf
	// This is Inf as bytes: this is hopefully unnecessary, but I don't know at this point.
	InfBytes []byte
	// Conn holds the configuration for connecting to the server.
	ConnCfg
	// Cfg holds the client configuration (how the client behaves).
	Cfg

	// Healthbeat buffers
	CPUData [][]byte
	MemData [][]byte

	muSend sync.Mutex
	WS     *websocket.Conn
	// Channel for outbound binary messages.  The message is assumed to be a
	// websocket.Binary type
	SendB       chan []byte
	SendStr     chan string
	mu          sync.Mutex
	isConnected bool
	ServerURL   url.URL
}

func New(id uint32, name string) *Client {
	bldr := flatbuffers.NewBuilder(0)
	n := bldr.CreateString(name)
	InfStart(bldr)
	InfAddID(bldr, id)
	InfAddHostname(bldr, n)
	bldr.Finish(InfEnd(bldr))
	return &Client{
		Inf:      GetRootAsInf(bldr.Bytes[bldr.Head():], 0),
		InfBytes: bldr.Bytes[bldr.Head():],
		// A really small buffer:
		// TODO: rethink this vis-a-vis what happens when recipient isn't there
		// or if it goes away during sending and possibly caching items to be sent.
		SendB:   make(chan []byte, 10),
		SendStr: make(chan string, 10),
	}
}

// Connect handles connecting to the server and returns the connection status.
// The client will attempt to connect until it has either succeeded or the
// connection retry period has been exceeded.  A retry is done every 5 seconds.
//
// If the client is already connected, nothing will be done.
func (c *Client) Connect() bool {
	// If already connected, return that fact.
	if c.IsConnected() {
		return true
	}
	start := time.Now()
	retryEnd := start.Add(c.ConnectPeriod)
	// connect to server; retry until the retry period has expired
	for {
		if time.Now().After(retryEnd) {
			fmt.Fprintf(os.Stderr, "timed out while trying to connect to the server: %s\n", c.ServerURL.String())
			return false
		}
		err := c.DialServer()
		if err == nil {
			break
		}
		time.Sleep(c.ConnectInterval)
		fmt.Printf("unable to connect to the server %s: retrying...\n", c.ServerURL.String())
	}
	// Send the ClientInf.  If the ID == 0 or it can't be found, the server will
	// respond with one.  Retry until the server responds, or until the
	// reconnectPeriod has expired.
	err := c.WS.WriteMessage(websocket.BinaryMessage, c.InfBytes)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while sending ID: %s\n", err)
		c.WS.Close()
		return false
	}

	typ, p, err := c.WS.ReadMessage()
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while Reading ID response: %s\n", err)
		c.WS.Close()
		return false
	}

	switch typ {
	case websocket.BinaryMessage:
		// a binary response is ClientInf
		c.Inf = GetRootAsInf(p, 0)
		c.InfBytes = p
		fmt.Printf("new ID: %d\n", c.Inf.ID())
	case websocket.TextMessage:
		fmt.Printf("%s\n", string(p))
	default:
		fmt.Printf("unexpected welcome response type %d: %v\n", typ, p)
		c.WS.Close()
		return false
	}
	// the next message is the current Cfg, which is applied to the clientID
	typ, p, err = c.WS.ReadMessage()
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while Reading ID response: %s\n", err)
		c.WS.Close()
		return false
	}

	switch typ {
	case websocket.BinaryMessage:
		// decode the message (it's in flatbuffer format), and apply to the cfg
		c.processBinaryMessage(p)
	default:
		fmt.Printf("unexpected message type from the server: was expecting the client's cfg: %d: %v\n", typ, p)
		fmt.Println("closing the connection and exiting")
		c.WS.Close()
		return false
	}

	c.mu.Lock()
	c.isConnected = true
	c.mu.Unlock()
	return true
}

func (c *Client) DialServer() error {
	var err error
	c.WS, _, err = websocket.DefaultDialer.Dial(c.ServerURL.String(), nil)
	return err
}

func (c *Client) MessageWriter(doneCh chan struct{}) {
	pingPeriod := time.Duration(c.Cfg.PingPeriod())
	defer close(doneCh)
	for {
		select {
		case p, ok := <-c.SendB:
			// don't send if not connected
			if !c.IsConnected() {
				continue
			}
			if !ok {
				c.WS.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			err := c.WS.WriteMessage(websocket.BinaryMessage, p)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error writing message: %s\n", err)
			}
		case <-time.After(pingPeriod):
			// only ping if we are connected
			if !c.IsConnected() {
				continue
			}
			err := c.WS.WriteMessage(websocket.PingMessage, []byte("ping"))
			if err != nil {
				fmt.Fprintf(os.Stderr, "ping error: %s\n", err)
			}
		}
	}
}

func (c *Client) Reconnect() bool {
	c.mu.Lock()
	c.isConnected = false
	c.mu.Unlock()
	for i := 0; i < 4; i++ {
		b := c.Connect()
		if b {
			fmt.Println("reconnect true")
			return b
		}
	}
	return false
}

func (c *Client) Listen(doneCh chan struct{}) {
	// loop until there's a done signal
	defer close(doneCh)
	for {
		typ, p, err := c.WS.ReadMessage()
		if err != nil {
			fmt.Fprintf(os.Stderr, "error reading message: %s\n", err)
			if _, ok := err.(*websocket.CloseError); !ok {
				return
			}
			fmt.Println("reconnecting from read messages")
			connected := c.Reconnect()
			if connected {
				continue
			}
			fmt.Fprint(os.Stderr, "unable to reconnect to server")
			return
		}
		switch typ {
		case websocket.TextMessage:
			fmt.Printf("textmessage: %s\n", p)
			if bytes.Equal(p, autofact.AckMsg) {
				// if this is an acknowledgement message, do nothing
				continue
			}
			err := c.WS.WriteMessage(websocket.TextMessage, autofact.AckMsg)
			if err != nil {
				if _, ok := err.(*websocket.CloseError); !ok {
					return
				}
				fmt.Println("reconnect from writing message: textmessage")
				connected := c.Reconnect()
				if connected {
					continue
				}
				fmt.Fprint(os.Stderr, "unable to reconnect to server")
				return
			}
		case websocket.BinaryMessage:
			err = c.WS.WriteMessage(websocket.TextMessage, autofact.AckMsg)
			if err != nil {
				fmt.Fprintf(os.Stderr, "error writing binary message: %s\n", err)
				if _, ok := err.(*websocket.CloseError); !ok {
					return
				}
				fmt.Println("reconnect from writing message: binarymessage")
				connected := c.Reconnect()
				if connected {
					continue
				}
				fmt.Fprint(os.Stderr, "unable to reconnect to server")
				return
			}
			c.processBinaryMessage(p)
		case websocket.CloseMessage:
			fmt.Printf("closemessage: %x\n", p)
			fmt.Println("reconnect from writing message: closemessage")
			connected := c.Reconnect()
			if connected {
				continue
			}
			fmt.Fprint(os.Stderr, "unable to reconnect to server")
			return
		}
	}
}

// IsConnected returns if the client is connected.
func (c *Client) IsConnected() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.isConnected
}

func (c *Client) PingHandler(msg string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	fmt.Printf("ping: %s\n", msg)
	return c.WS.WriteMessage(websocket.PongMessage, []byte("ping"))
}

func (c *Client) PongHandler(msg string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	fmt.Printf("pong: %s\n", msg)
	return c.WS.WriteMessage(websocket.PingMessage, []byte("pong"))
}

// TODO: should CPUData be enclosed in a struct for locking purposes?
func (c *Client) EnqueueCPUData(data []byte) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.CPUData = append(c.CPUData, data)
	return len(c.CPUData)
}

// EnqueueMemData adds the received data to the MemData buffer.  The current
// number of entries in the buffer is returned.
func (c *Client) EnqueueMemData(data []byte) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.MemData = append(c.MemData, data)
	return len(c.MemData)
}

// If the message send fails, whatever was cached will be lost.
// TODO: the current stats get copied for send; should the stat slice
// get reset now?  There is possible data loss this way; but there's
// possible data loss if the stat slice gets appended to between this
// copy op and the send completing, unless a lock is held the entire time,
// which could block the stats reading process leading to data loss due to
// missed reads.  I'm thinking COW or delete now; but punting becasue it
// really doesn't matter as this is just an experiment, right now.
// This also applies to CPUDatasFB
// TODO: should the messages to be sent be copied to a send cache so that
// there isn't data loss on a failed send?  Consecutive PushPeriods that
// failed to send may be problematic in that situation.
func (c *Client) FlushCPUData() [][]byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	data := make([][]byte, len(c.CPUData))
	copy(data, c.CPUData)
	c.CPUData = nil
	return data
}

// FlushMemData returns a copy of the client's MemData and nils the client's
// cache.
// The notes above apply here too.
func (c *Client) FlushMemData() [][]byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	data := make([][]byte, len(c.MemData))
	copy(data, c.MemData)
	c.MemData = nil
	return data
}

func (c *Client) Healthbeat() {
	// An interval of 0 means no healthbeat
	if c.Cfg.HealthbeatInterval() == 0 {
		return
	}
	cpuCh := make(chan []byte)
	memCh := make(chan []byte)
	go sysinfo.CPUDataTicker(time.Duration(c.Cfg.HealthbeatInterval()), cpuCh)
	go sysinfo.MemDataTicker(time.Duration(c.Cfg.HealthbeatInterval()), memCh)
	t := time.NewTicker(time.Duration(c.Cfg.HealthbeatPushPeriod()))
	defer t.Stop()
	for {
		select {
		case data, ok := <-cpuCh:
			if !ok {
				fmt.Println("cpu stats chan closed")
				goto done
			}
			c.EnqueueCPUData(data)
		case data, ok := <-memCh:
			if !ok {
				fmt.Println("cpu stats chan closed")
				goto done
			}
			c.EnqueueMemData(data)
		case <-t.C:
			if !c.IsConnected() {
				continue
			}
			c.SendData(message.CPUData, c.FlushCPUData())
			c.SendData(message.MemData, c.FlushMemData())
		}
	}
done:
	// Flush the buffer.
	c.SendData(message.CPUData, c.FlushCPUData())
	c.SendData(message.MemData, c.FlushMemData())
}

// SendData sends the received data to the server.  The caller checks to see
// if the client is connected to the server before calling.  If the connection
// is lost during processing, the cached stats will be lost.
// TODO:  should this be more resilient?
func (c *Client) SendData(kind message.Kind, data [][]byte) error {
	// for each data, send a message
	for _, v := range data {
		c.SendB <- message.Serialize(c.Inf.ID(), kind, v)
	}
	return nil
}

// SendMessage sends a single serialized message of type Kind.
func (c *Client) SendMessage(kind message.Kind, p []byte) {
	c.SendB <- message.Serialize(c.Inf.ID(), kind, p)
}

// binary messages are expected to be flatbuffer encoding of message.Message.
func (c *Client) processBinaryMessage(p []byte) error {
	// unmarshal the message
	msg := message.GetRootAsMessage(p, 0)
	// process according to kind
	k := message.Kind(msg.Kind())
	switch k {
	case message.ClientCfg:
		c.Cfg.Deserialize(msg.DataBytes())
	default:
		fmt.Println("unknown message kind")
		fmt.Println(string(p))
	}
	return nil
}

// Serialize serializes the Inf using flatbuffers and returns the []byte.
func (i *Inf) Serialize() []byte {
	bldr := flatbuffers.NewBuilder(0)
	h := bldr.CreateByteString(i.Hostname())
	r := bldr.CreateByteString(i.Region())
	z := bldr.CreateByteString(i.Zone())
	d := bldr.CreateByteString(i.DC())
	InfStart(bldr)
	InfAddID(bldr, i.ID())
	InfAddHostname(bldr, h)
	InfAddRegion(bldr, r)
	InfAddZone(bldr, z)
	InfAddDC(bldr, d)
	bldr.Finish(InfEnd(bldr))
	return bldr.Bytes[bldr.Head():]
}
<|endoftext|>


<commit_before>// The Toxiproxy Go client provides a wrapper around the Toxiproxy HTTP API for
// testing the resiliency of Go applications.
package client

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
)

// Client holds information about where to connect to Toxiproxy.
type Client struct {
	endpoint string
}

type Fields map[string]interface{}

// Proxy represents a Proxy.
type Proxy struct {
	Name     string `json:"name"`     // The name of the proxy
	Listen   string `json:"listen"`   // The address the proxy listens on
	Upstream string `json:"upstream"` // The upstream address to proxy to
	Enabled  bool   `json:"enabled"`  // Whether the proxy is enabled

	ToxicsUpstream   map[string]interface{} `json:"upstream_toxics"`   // Toxics in the upstream direction
	ToxicsDownstream map[string]interface{} `json:"downstream_toxics"` // Toxics in the downstream direction

	client *Client
}

// NewClient creates a new client which provides the base of all communication
// with Toxiproxy. Endpoint is the address to the proxy (e.g. localhost:8474 if
// not overriden)
func NewClient(endpoint string) *Client {
	return &Client{endpoint: endpoint}
}

// Proxies returns a map with all the proxies.
func (client *Client) Proxies() (map[string]*Proxy, error) {
	resp, err := http.Get(client.endpoint + "/proxies")
	if err != nil {
		return nil, err
	}

	proxies := make(map[string]*Proxy)
	err = json.NewDecoder(resp.Body).Decode(&proxies)
	if err != nil {
		return nil, err
	}
	for _, proxy := range proxies {
		proxy.client = client
	}

	return proxies, nil
}

// NewProxy instantiates a new proxy instance. Note Create() must be called on
// it to create it. The Enabled field must be set to true, otherwise the Proxy
// will not be enabled when created.
func (client *Client) NewProxy(proxy *Proxy) *Proxy {
	if proxy == nil {
		proxy = &Proxy{}
	}

	proxy.client = client
	return proxy
}

// Create creates a new proxy.
func (proxy *Proxy) Create() error {
	request, err := json.Marshal(proxy)
	if err != nil {
		return err
	}

	resp, err := http.Post(proxy.client.endpoint+"/proxies", "application/json", bytes.NewReader(request))
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusCreated {
		// TODO  better error
		return fmt.Errorf("omg error code %d", resp.StatusCode)
	}

	proxy = new(Proxy)
	err = json.NewDecoder(resp.Body).Decode(&proxy)
	if err != nil {
		return err
	}

	return nil
}

// Proxy returns a proxy by name.
func (client *Client) Proxy(name string) (*Proxy, error) {
	// TODO url encode
	resp, err := http.Get(client.endpoint + "/proxies/" + name)
	if err != nil {
		return nil, err
	}

	proxy := client.NewProxy(nil)
	err = json.NewDecoder(resp.Body).Decode(proxy)
	if err != nil {
		return nil, err
	}

	return proxy, nil
}

// Save saves changes to a proxy such as its enabled status.
func (proxy *Proxy) Save() error {
	request, err := json.Marshal(proxy)
	if err != nil {
		return err
	}

	resp, err := http.Post(proxy.client.endpoint+"/proxies/"+proxy.Name, "application/json", bytes.NewReader(request))
	if err != nil {
		return err
	}

	err = json.NewDecoder(resp.Body).Decode(proxy)
	if err != nil {
		return err
	}

	return nil
}

// Delete a proxy which will cause it to stop listening and delete all
// information associated with it. If you just wish to stop and later enable a
// proxy, set the `Enabled` field to `false` and call `Save()`.
func (proxy *Proxy) Delete() error {
	httpClient := &http.Client{}
	req, err := http.NewRequest("DELETE", proxy.client.endpoint+"/proxies/"+proxy.Name, nil)

	if err != nil {
		return err
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusNoContent {
		// TODO better error
		return errors.New("Status code bad")
	}

	return nil
}

// Toxics returns a map of all the toxics and their attributes for a direction.
func (proxy *Proxy) Toxics(direction string) (map[string]interface{}, error) {
	resp, err := http.Get(proxy.client.endpoint + "/proxies/" + proxy.Name + "/" + direction + "/toxics")
	if err != nil {
		return nil, err
	}

	toxics := make(map[string]interface{})
	err = json.NewDecoder(resp.Body).Decode(&toxics)
	if err != nil {
		return nil, err
	}

	return toxics, nil
}

// SetToxic sets the parameters for a toxic with a given name in the direction.
// See https://github.com/Shopify/toxiproxy#toxics for a list of all Toxics.
func (proxy *Proxy) SetToxic(name string, direction string, fields Fields) (map[string]interface{}, error) {
	request, err := json.Marshal(fields)
	if err != nil {
		return nil, err
	}

	resp, err := http.Post(proxy.client.endpoint+"/proxies/"+proxy.Name+"/"+direction+"/toxics/"+name, "application/json", bytes.NewReader(request))
	if err != nil {
		return nil, err
	}

	toxics := make(map[string]interface{})
	err = json.NewDecoder(resp.Body).Decode(&toxics)
	if err != nil {
		return nil, err
	}

	return toxics, nil
}

// Toxics lists all proxies and toxics.
func (client *Client) Toxics() (map[string]*Proxy, error) {
	resp, err := http.Get(client.endpoint + "/toxics")
	if err != nil {
		return nil, err
	}

	proxies := make(map[string]*Proxy)
	err = json.NewDecoder(resp.Body).Decode(&proxies)
	if err != nil {
		return nil, err
	}

	return proxies, nil
}

// ResetState resets the state of all proxies and toxics in Toxiproxy.
func (client *Client) ResetState() error {
	resp, err := http.Get(client.endpoint + "/reset")
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusNoContent {
		// TODO better error
		return errors.New("unable to reset")
	}

	return nil
}
<commit_msg>client: fix package level comment<commit_after>// Package client provides a Toxiproxy client provides a wrapper around the
// Toxiproxy HTTP API for testing the resiliency of Go applications.
package client

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
)

// Client holds information about where to connect to Toxiproxy.
type Client struct {
	endpoint string
}

type Fields map[string]interface{}

// Proxy represents a Proxy.
type Proxy struct {
	Name     string `json:"name"`     // The name of the proxy
	Listen   string `json:"listen"`   // The address the proxy listens on
	Upstream string `json:"upstream"` // The upstream address to proxy to
	Enabled  bool   `json:"enabled"`  // Whether the proxy is enabled

	ToxicsUpstream   map[string]interface{} `json:"upstream_toxics"`   // Toxics in the upstream direction
	ToxicsDownstream map[string]interface{} `json:"downstream_toxics"` // Toxics in the downstream direction

	client *Client
}

// NewClient creates a new client which provides the base of all communication
// with Toxiproxy. Endpoint is the address to the proxy (e.g. localhost:8474 if
// not overriden)
func NewClient(endpoint string) *Client {
	return &Client{endpoint: endpoint}
}

// Proxies returns a map with all the proxies.
func (client *Client) Proxies() (map[string]*Proxy, error) {
	resp, err := http.Get(client.endpoint + "/proxies")
	if err != nil {
		return nil, err
	}

	proxies := make(map[string]*Proxy)
	err = json.NewDecoder(resp.Body).Decode(&proxies)
	if err != nil {
		return nil, err
	}
	for _, proxy := range proxies {
		proxy.client = client
	}

	return proxies, nil
}

// NewProxy instantiates a new proxy instance. Note Create() must be called on
// it to create it. The Enabled field must be set to true, otherwise the Proxy
// will not be enabled when created.
func (client *Client) NewProxy(proxy *Proxy) *Proxy {
	if proxy == nil {
		proxy = &Proxy{}
	}

	proxy.client = client
	return proxy
}

// Create creates a new proxy.
func (proxy *Proxy) Create() error {
	request, err := json.Marshal(proxy)
	if err != nil {
		return err
	}

	resp, err := http.Post(proxy.client.endpoint+"/proxies", "application/json", bytes.NewReader(request))
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusCreated {
		// TODO  better error
		return fmt.Errorf("omg error code %d", resp.StatusCode)
	}

	proxy = new(Proxy)
	err = json.NewDecoder(resp.Body).Decode(&proxy)
	if err != nil {
		return err
	}

	return nil
}

// Proxy returns a proxy by name.
func (client *Client) Proxy(name string) (*Proxy, error) {
	// TODO url encode
	resp, err := http.Get(client.endpoint + "/proxies/" + name)
	if err != nil {
		return nil, err
	}

	proxy := client.NewProxy(nil)
	err = json.NewDecoder(resp.Body).Decode(proxy)
	if err != nil {
		return nil, err
	}

	return proxy, nil
}

// Save saves changes to a proxy such as its enabled status.
func (proxy *Proxy) Save() error {
	request, err := json.Marshal(proxy)
	if err != nil {
		return err
	}

	resp, err := http.Post(proxy.client.endpoint+"/proxies/"+proxy.Name, "application/json", bytes.NewReader(request))
	if err != nil {
		return err
	}

	err = json.NewDecoder(resp.Body).Decode(proxy)
	if err != nil {
		return err
	}

	return nil
}

// Delete a proxy which will cause it to stop listening and delete all
// information associated with it. If you just wish to stop and later enable a
// proxy, set the `Enabled` field to `false` and call `Save()`.
func (proxy *Proxy) Delete() error {
	httpClient := &http.Client{}
	req, err := http.NewRequest("DELETE", proxy.client.endpoint+"/proxies/"+proxy.Name, nil)

	if err != nil {
		return err
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusNoContent {
		// TODO better error
		return errors.New("Status code bad")
	}

	return nil
}

// Toxics returns a map of all the toxics and their attributes for a direction.
func (proxy *Proxy) Toxics(direction string) (map[string]interface{}, error) {
	resp, err := http.Get(proxy.client.endpoint + "/proxies/" + proxy.Name + "/" + direction + "/toxics")
	if err != nil {
		return nil, err
	}

	toxics := make(map[string]interface{})
	err = json.NewDecoder(resp.Body).Decode(&toxics)
	if err != nil {
		return nil, err
	}

	return toxics, nil
}

// SetToxic sets the parameters for a toxic with a given name in the direction.
// See https://github.com/Shopify/toxiproxy#toxics for a list of all Toxics.
func (proxy *Proxy) SetToxic(name string, direction string, fields Fields) (map[string]interface{}, error) {
	request, err := json.Marshal(fields)
	if err != nil {
		return nil, err
	}

	resp, err := http.Post(proxy.client.endpoint+"/proxies/"+proxy.Name+"/"+direction+"/toxics/"+name, "application/json", bytes.NewReader(request))
	if err != nil {
		return nil, err
	}

	toxics := make(map[string]interface{})
	err = json.NewDecoder(resp.Body).Decode(&toxics)
	if err != nil {
		return nil, err
	}

	return toxics, nil
}

// Toxics lists all proxies and toxics.
func (client *Client) Toxics() (map[string]*Proxy, error) {
	resp, err := http.Get(client.endpoint + "/toxics")
	if err != nil {
		return nil, err
	}

	proxies := make(map[string]*Proxy)
	err = json.NewDecoder(resp.Body).Decode(&proxies)
	if err != nil {
		return nil, err
	}

	return proxies, nil
}

// ResetState resets the state of all proxies and toxics in Toxiproxy.
func (client *Client) ResetState() error {
	resp, err := http.Get(client.endpoint + "/reset")
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusNoContent {
		// TODO better error
		return errors.New("unable to reset")
	}

	return nil
}
<|endoftext|>


<commit_before>package gitmediaclient

import (
	".."
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/cheggaaa/pb"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
)

const (
	gitMediaType     = "application/vnd.git-media"
	gitMediaMetaType = gitMediaType + "+json; charset=utf-8"
	gitMediaHeader   = "0014 git media v1\n"
)

func Options(filehash string) error {
	oid := filepath.Base(filehash)
	_, err := os.Stat(filehash)
	if err != nil {
		return err
	}

	req, creds, err := clientRequest("OPTIONS", oid)
	if err != nil {
		return err
	}

	_, err = doRequest(req, creds)
	if err != nil {
		return err
	}

	return nil
}

func Put(filehash, filename string) error {
	if filename == "" {
		filename = filehash
	}

	oid := filepath.Base(filehash)
	stat, err := os.Stat(filehash)
	if err != nil {
		return err
	}

	file, err := os.Open(filehash)
	if err != nil {
		return err
	}

	req, creds, err := clientRequest("PUT", oid)
	if err != nil {
		return err
	}

	bar := pb.StartNew(int(stat.Size()))
	bar.SetUnits(pb.U_BYTES)
	bar.Start()

	req.Header.Set("Content-Type", gitMediaType)
	req.Header.Set("Accept", gitMediaMetaType)
	req.Body = ioutil.NopCloser(bar.NewProxyReader(file))
	req.ContentLength = stat.Size()

	fmt.Printf("Sending %s\n", filename)

	_, err = doRequest(req, creds)
	if err != nil {
		return err
	}

	return nil
}

func Get(filename string) (io.ReadCloser, error) {
	oid := filepath.Base(filename)
	if stat, err := os.Stat(filename); err != nil || stat == nil {
		req, creds, err := clientRequest("GET", oid)
		if err != nil {
			return nil, err
		}

		req.Header.Set("Accept", gitMediaType)
		res, err := doRequest(req, creds)

		if err != nil {
			return nil, err
		}

		header := make([]byte, len(gitMediaHeader))
		_, err = io.ReadAtLeast(res.Body, header, len(gitMediaHeader))
		if err != nil {
			return nil, err
		}

		if string(header) != gitMediaHeader {
			return nil, errors.New("Invalid header")
		}

		return res.Body, nil
	}

	return os.Open(filename)
}

func doRequest(req *http.Request, creds Creds) (*http.Response, error) {
	res, err := http.DefaultClient.Do(req)

	if err == nil {
		if res.StatusCode > 299 {
			execCreds(creds, "reject")

			apierr := &Error{}
			dec := json.NewDecoder(res.Body)
			if err := dec.Decode(apierr); err != nil {
				return res, err
			}

			return res, apierr
		}
	} else {
		execCreds(creds, "approve")
	}

	return res, err
}

func clientRequest(method, oid string) (*http.Request, Creds, error) {
	u := ObjectUrl(oid)
	req, err := http.NewRequest(method, u.String(), nil)
	if err == nil {
		creds, err := credentials(u)
		if err != nil {
			return req, nil, err
		}

		token := fmt.Sprintf("%s:%s", creds["username"], creds["password"])
		auth := "Basic " + base64.URLEncoding.EncodeToString([]byte(token))
		req.Header.Set("Authorization", auth)
		return req, creds, nil
	}

	return req, nil, err
}

func ObjectUrl(oid string) *url.URL {
	c := gitmedia.Config
	u, _ := url.Parse(c.Endpoint())
	u.Path = filepath.Join(u.Path, "/objects/"+oid)
	return u
}

type Error struct {
	Message   string `json:"message"`
	RequestId string `json:"request_id,omitempty"`
}

func (e *Error) Error() string {
	return e.Message
}
<commit_msg>ンンン ンンン ンン<commit_after>package gitmediaclient

import (
	".."
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/cheggaaa/pb"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
)

const (
	gitMediaType      = "application/vnd.git-media"
	gitMediaMetaType  = gitMediaType + "+json; charset=utf-8"
	gitMediaHeader    = "--git-media."
	gitBoundaryLength = 61
)

func Options(filehash string) error {
	oid := filepath.Base(filehash)
	_, err := os.Stat(filehash)
	if err != nil {
		return err
	}

	req, creds, err := clientRequest("OPTIONS", oid)
	if err != nil {
		return err
	}

	_, err = doRequest(req, creds)
	if err != nil {
		return err
	}

	return nil
}

func Put(filehash, filename string) error {
	if filename == "" {
		filename = filehash
	}

	oid := filepath.Base(filehash)
	stat, err := os.Stat(filehash)
	if err != nil {
		return err
	}

	file, err := os.Open(filehash)
	if err != nil {
		return err
	}

	req, creds, err := clientRequest("PUT", oid)
	if err != nil {
		return err
	}

	bar := pb.StartNew(int(stat.Size()))
	bar.SetUnits(pb.U_BYTES)
	bar.Start()

	req.Header.Set("Content-Type", gitMediaType)
	req.Header.Set("Accept", gitMediaMetaType)
	req.Body = ioutil.NopCloser(bar.NewProxyReader(file))
	req.ContentLength = stat.Size()

	fmt.Printf("Sending %s\n", filename)

	_, err = doRequest(req, creds)
	if err != nil {
		return err
	}

	return nil
}

func Get(filename string) (io.ReadCloser, error) {
	oid := filepath.Base(filename)
	if stat, err := os.Stat(filename); err != nil || stat == nil {
		req, creds, err := clientRequest("GET", oid)
		if err != nil {
			return nil, err
		}

		req.Header.Set("Accept", gitMediaType)
		res, err := doRequest(req, creds)

		if err != nil {
			return nil, err
		}

		header := make([]byte, len(gitMediaHeader)+gitBoundaryLength)
		_, err = io.ReadAtLeast(res.Body, header, len(gitMediaHeader)+gitBoundaryLength)
		if err != nil {
			return nil, err
		}

		if !strings.HasPrefix(string(header), gitMediaHeader) {
			return nil, errors.New("Invalid header")
		}

		return res.Body, nil
	}

	return os.Open(filename)
}

func doRequest(req *http.Request, creds Creds) (*http.Response, error) {
	res, err := http.DefaultClient.Do(req)

	if err == nil {
		if res.StatusCode > 299 {
			execCreds(creds, "reject")

			apierr := &Error{}
			dec := json.NewDecoder(res.Body)
			if err := dec.Decode(apierr); err != nil {
				return res, err
			}

			return res, apierr
		}
	} else {
		execCreds(creds, "approve")
	}

	return res, err
}

func clientRequest(method, oid string) (*http.Request, Creds, error) {
	u := ObjectUrl(oid)
	req, err := http.NewRequest(method, u.String(), nil)
	if err == nil {
		creds, err := credentials(u)
		if err != nil {
			return req, nil, err
		}

		token := fmt.Sprintf("%s:%s", creds["username"], creds["password"])
		auth := "Basic " + base64.URLEncoding.EncodeToString([]byte(token))
		req.Header.Set("Authorization", auth)
		return req, creds, nil
	}

	return req, nil, err
}

func ObjectUrl(oid string) *url.URL {
	c := gitmedia.Config
	u, _ := url.Parse(c.Endpoint())
	u.Path = filepath.Join(u.Path, "/objects/"+oid)
	return u
}

type Error struct {
	Message   string `json:"message"`
	RequestId string `json:"request_id,omitempty"`
}

func (e *Error) Error() string {
	return e.Message
}
<|endoftext|>