<commit_before>/* Copyright (C) 2014 CompleteDB LLC.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the Apache License Version 2.0 http://www.apache.org/licenses.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 */

package pubsubsql

import (
	"container/list"
	"encoding/json"
	"net"
	"time"
)

var _CLIENT_DEFAULT_BUFFER_SIZE int = 2048

// respnoseData holds unmarshaled result from pubsubsql JSON response
type responseData struct {
	Status   string
	Msg      string
	Action   string
	Id       string
	PubSubId string
	Rows     int
	Fromrow  int
	Torow    int
	Columns  []string
	Data     [][]string
}

func (this *responseData) reset() {
	this.Status = ""
	this.Msg = ""
	this.Action = ""
	this.PubSubId = ""
	this.Rows = 0
	this.Fromrow = 0
	this.Torow = 0
	this.Columns = nil
	this.Data = nil
}

type Client struct {
	address   string
	rw        netHelper
	requestId uint32
	err       string
	rawjson   []byte
	//
	response responseData
	record   int
	columns  map[string]int

	// pubsub back log
	backlog list.List
}

//Connect connects the Client to the pubsubsql server.
//Address string has the form host:port.
func (this *Client) Connect(address string) bool {
	this.address = address
	this.Disconnect()
	conn, err := net.DialTimeout("tcp", this.address, time.Millisecond*1000)
	if err != nil {
		this.setError(err)
		return false
	}
	this.rw.set(conn, _CLIENT_DEFAULT_BUFFER_SIZE)
	return true
}

//Disconnect disconnects the Client from the pubsubsql server.
func (this *Client) Disconnect() {
	this.write("close")
	// write may generate error so we reset after instead
	this.reset()
	this.rw.close()
}

//Connected returns true if the Client is currently connected to the pubsubsql server.
func (this *Client) Connected() bool {
	return this.rw.valid()
}

//Ok determines if the last command executed against the pubsubsql server succeeded.
func (this *Client) Ok() bool {
	return this.err == ""
}

//Failed determines if the last command executed against the pubsubsql server failed.
func (this *Client) Failed() bool {
	return !this.Ok()
}

//Error returns an error message when the last command executed against
//the pubsubsql server fails.

//Functions that may generate an error are [Connect, Execute, NextRow, WaitForPubSub]
func (this *Client) Error() string {
	return this.err
}

//Execute executes a command against the pubsubsql server and returns true on success.
//The pubsubsql server returns to the Client a response in JSON format.
func (this *Client) Execute(command string) bool {
	this.reset()
	ok := this.write(command)
	var bytes []byte
	var header *netHeader
	for ok {
		this.reset()
		header, bytes, ok = this.read()
		if !ok {
			break
		}
		if header.RequestId == this.requestId {
			// response we are waiting for
			return this.unmarshalJSON(bytes)
		} else if header.RequestId == 0 {
			// pubsub action, save it and skip it for now
			// will be proccesed next time WaitPubSub is called
			//WE MUST COPY BYTES SINCE THEY ARE REUSED IN NetHelper
			t := make([]byte, header.MessageSize, header.MessageSize)
			copy(t, bytes[0:header.MessageSize])
			this.backlog.PushBack(t)
		} else if header.RequestId < this.requestId {
			// we did not read full result set from previous command ignore it or report error?
			// for now lets ignore it, continue reading until we hit our request id
			this.reset()
		} else {
			// this should never happen
			this.setErrorString("protocol error invalid requestId")
			ok = false
		}
	}
	return ok
}

//Stream sends a command against the pubsubsql server and returns true on success.
//The pubsubsql server does not return a response to the Client.
func (this *Client) Stream(command string) bool {
	this.reset()
	//TODO optimize
	return this.write("stream " + command)
}

//JSON returns a response string in JSON format from the
//last command executed against the pubsubsql server.
func (this *Client) JSON() string {
	return string(this.rawjson)
}

//Action returns an action string from the response
//returned by the last command executed against the pubsubsql server.
//Valid actions are [status, insert, select, delete, update, add, remove, subscribe, unsubscribe]
func (this *Client) Action() string {
	return this.response.Action
}

//PubSubId returns a unique identifier generated by the pubsubsql server when
//a Client subscribes to a table. If the client has subscribed to more than  one table,
//PubSubId should be used by the Client to uniquely identify messages
//published by the pubsubsql server.
func (this *Client) PubSubId() string {
	return this.response.PubSubId
}

//RowCount returns the number of rows in the result set returned by the pubsubsql server.
func (this *Client) RowCount() int {
	return this.response.Rows
}

//NextRow is used to move to the next row in the result set returned by the pubsubsql server.
//When called for the first time, NextRow moves to the first row in the result set.
//Returns false when all rows are read or if there is an error.
//To find out if false was returned because of an error, use Ok or Failed functions.
func (this *Client) NextRow() bool {
	for this.Ok() {
		// no result set
		if this.response.Rows == 0 {
			return false
		}
		if this.response.Fromrow == 0 || this.response.Torow == 0 {
			return false
		}
		// the current record is valid
		this.record++
		if this.record <= (this.response.Torow - this.response.Fromrow) {
			return true
		}
		// we reached the end of result set
		if this.response.Rows == this.response.Torow {
			// gaurd against over fill
			this.record--
			return false
		}
		// if we are here there is another batch
		this.reset()
		header, bytes, ok := this.read()
		if !ok {
			return false
		}
		// should not happen but check anyway
		// when RequestId is 0 it means we are reading published data
		if header.RequestId > 0 && header.RequestId != this.requestId {
			this.setErrorString("protocol error")
			return false
		}
		// we got another batch unmarshall the data
		this.unmarshalJSON(bytes)
	}
	return false
}

//Value returns the value within the current row for the given column name.
//If the column name does not exist, Value returns an empty string.
func (this *Client) Value(column string) string {
	ordinal, ok := this.columns[column]
	if !ok {
		return ""
	}
	return this.ValueByOrdinal(ordinal)
}

//ValueByOrdinal returns the value within the current row for the given column ordinal.
//The column ordinal represents the zero based position of the column in the Columns collection of the result set.
//If the column ordinal is out of range, ValueByOrdinal returns an empty string.
func (this *Client) ValueByOrdinal(ordinal int) string {
	if this.record < 0 || this.record >= len(this.response.Data) {
		return ""
	}
	if ordinal >= len(this.response.Data[this.record]) {
		return ""
	}
	return this.response.Data[this.record][ordinal]
}

//HasColumn determines if the column name exists in the columns collection of the result set.
func (this *Client) HasColumn(column string) bool {
	_, ok := this.columns[column]
	return ok
}

//ColumnCount returns the number of columns in the columns collection of the result set.
func (this *Client) ColumnCount() int {
	return len(this.response.Columns)
}

//Columns returns the column names in the columns collection of the result set.
func (this *Client) Columns() []string {
	return this.response.Columns
}

//WaitForPubSub waits until the pubsubsql server publishes a message for
// the subscribed Client or until the timeout interval elapses.
//Returns false when timeout interval elapses or if there is and error.
//To find out if false was returned because of an error, use Ok or Failed functions.
func (this *Client) WaitForPubSub(timeout int) bool {
	var bytes []byte
	for {
		this.reset()
		// process backlog first
		bytes = this.popBacklog()
		if len(bytes) > 0 {
			return this.unmarshalJSON(bytes)
		}
		header, temp, success, timedout := this.readTimeout(int64(timeout))
		bytes = temp
		if !success || timedout {
			return false
		}
		if header.RequestId == 0 {
			return this.unmarshalJSON(bytes)
		}
		// this is not pubsub message; are we reading abandoned cursor?
		// ignore and keep trying
	}
	return false
}

func (this *Client) popBacklog() []byte {
	element := this.backlog.Front()
	if element != nil {
		bytes := element.Value.([]byte)
		this.backlog.Remove(element)
		return bytes
	}
	return nil
}

func (this *Client) unmarshalJSON(bytes []byte) bool {
	this.rawjson = bytes
	err := json.Unmarshal(bytes, &this.response)
	if err != nil {
		this.setError(err)
		return false
	}
	if this.response.Status != "ok" {
		this.setErrorString(this.response.Msg)
		return false
	}
	this.setColumns()
	return true
}

func (this *Client) setColumns() {
	if len(this.response.Columns) == 0 {
		return
	}
	this.columns = make(map[string]int, cap(this.response.Columns))
	for ordinal, column := range this.response.Columns {
		this.columns[column] = ordinal
	}
}

func (this *Client) reset() {
	this.resetError()
	this.response.reset()
	this.rawjson = nil
	this.record = -1
}

func (this *Client) resetError() {
	this.err = ""
}

func (this *Client) setErrorString(err string) {
	this.reset()
	this.err = err
}

func (this *Client) setError(err error) {
	this.setErrorString(err.Error())
}

func (this *Client) write(message string) bool {
	this.requestId++
	if !this.rw.valid() {
		this.setErrorString("Not connected")
		return false
	}
	err := this.rw.writeHeaderAndMessage(this.requestId, []byte(message))
	if err != nil {
		this.setError(err)
		return false
	}
	return true
}

func (this *Client) readTimeout(timeout int64) (*netHeader, []byte, bool, bool) {
	if !this.rw.valid() {
		this.setErrorString("Not connected")
		return nil, nil, false, false
	}
	header, bytes, err, timedout := this.rw.readMessageTimeout(timeout)
	if timedout {
		return nil, nil, true, true
	}
	// error
	if err != nil {
		this.setError(err)
		return nil, nil, false, false
	}
	// success
	return header, bytes, true, false

}

func (this *Client) read() (*netHeader, []byte, bool) {
	var MAX_READ_TIMEOUT_MILLISECONDS int64 = 1000 * 60 * 3
	header, bytes, success, timedout := this.readTimeout(MAX_READ_TIMEOUT_MILLISECONDS)
	if timedout {
		this.setErrorString("Read timed out")
		return nil, nil, false
	}
	return header, bytes, success
}
<commit_msg>client update doc<commit_after>/* Copyright (C) 2014 CompleteDB LLC.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the Apache License Version 2.0 http://www.apache.org/licenses.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 */

package pubsubsql

import (
	"container/list"
	"encoding/json"
	"net"
	"time"
)

var _CLIENT_DEFAULT_BUFFER_SIZE int = 2048

// respnoseData holds unmarshaled result from pubsubsql JSON response
type responseData struct {
	Status   string
	Msg      string
	Action   string
	Id       string
	PubSubId string
	Rows     int
	Fromrow  int
	Torow    int
	Columns  []string
	Data     [][]string
}

func (this *responseData) reset() {
	this.Status = ""
	this.Msg = ""
	this.Action = ""
	this.PubSubId = ""
	this.Rows = 0
	this.Fromrow = 0
	this.Torow = 0
	this.Columns = nil
	this.Data = nil
}

type Client struct {
	address   string
	rw        netHelper
	requestId uint32
	err       string
	rawjson   []byte
	//
	response responseData
	record   int
	columns  map[string]int

	// pubsub back log
	backlog list.List
}

//Connect connects the Client to the pubsubsql server.
//Address string has the form host:port.
func (this *Client) Connect(address string) bool {
	this.address = address
	this.Disconnect()
	conn, err := net.DialTimeout("tcp", this.address, time.Millisecond*1000)
	if err != nil {
		this.setError(err)
		return false
	}
	this.rw.set(conn, _CLIENT_DEFAULT_BUFFER_SIZE)
	return true
}

//Disconnect disconnects the Client from the pubsubsql server.
func (this *Client) Disconnect() {
	this.write("close")
	// write may generate error so we reset after instead
	this.reset()
	this.rw.close()
}

//Connected returns true if the Client is currently connected to the pubsubsql server.
func (this *Client) Connected() bool {
	return this.rw.valid()
}

//Ok determines if the last command executed against the pubsubsql server succeeded.
func (this *Client) Ok() bool {
	return this.err == ""
}

//Failed determines if the last command executed against the pubsubsql server failed.
func (this *Client) Failed() bool {
	return !this.Ok()
}

//Error returns an error message when the last command executed against
//the pubsubsql server fails.

//Functions that may generate an error are [Connect, Execute, NextRow, WaitForPubSub]
func (this *Client) Error() string {
	return this.err
}

//Execute executes a command against the pubsubsql server and returns true on success.
//The pubsubsql server returns to the Client a response in JSON format.
func (this *Client) Execute(command string) bool {
	this.reset()
	ok := this.write(command)
	var bytes []byte
	var header *netHeader
	for ok {
		this.reset()
		header, bytes, ok = this.read()
		if !ok {
			break
		}
		if header.RequestId == this.requestId {
			// response we are waiting for
			return this.unmarshalJSON(bytes)
		} else if header.RequestId == 0 {
			// pubsub action, save it and skip it for now
			// will be proccesed next time WaitPubSub is called
			//WE MUST COPY BYTES SINCE THEY ARE REUSED IN NetHelper
			t := make([]byte, header.MessageSize, header.MessageSize)
			copy(t, bytes[0:header.MessageSize])
			this.backlog.PushBack(t)
		} else if header.RequestId < this.requestId {
			// we did not read full result set from previous command ignore it or report error?
			// for now lets ignore it, continue reading until we hit our request id
			this.reset()
		} else {
			// this should never happen
			this.setErrorString("protocol error invalid requestId")
			ok = false
		}
	}
	return ok
}

//Stream sends a command to the pubsubsql server and returns true on success.
//The pubsubsql server does not return a response to the Client.
func (this *Client) Stream(command string) bool {
	this.reset()
	//TODO optimize
	return this.write("stream " + command)
}

//JSON returns a response string in JSON format from the
//last command executed against the pubsubsql server.
func (this *Client) JSON() string {
	return string(this.rawjson)
}

//Action returns an action string from the response
//returned by the last command executed against the pubsubsql server.
//Valid actions are [status, insert, select, delete, update, add, remove, subscribe, unsubscribe]
func (this *Client) Action() string {
	return this.response.Action
}

//PubSubId returns a unique identifier generated by the pubsubsql server when
//a Client subscribes to a table. If the client has subscribed to more than  one table,
//PubSubId should be used by the Client to uniquely identify messages
//published by the pubsubsql server.
func (this *Client) PubSubId() string {
	return this.response.PubSubId
}

//RowCount returns the number of rows in the result set returned by the pubsubsql server.
func (this *Client) RowCount() int {
	return this.response.Rows
}

//NextRow is used to move to the next row in the result set returned by the pubsubsql server.
//When called for the first time, NextRow moves to the first row in the result set.
//Returns false when all rows are read or if there is an error.
//To find out if false was returned because of an error, use Ok or Failed functions.
func (this *Client) NextRow() bool {
	for this.Ok() {
		// no result set
		if this.response.Rows == 0 {
			return false
		}
		if this.response.Fromrow == 0 || this.response.Torow == 0 {
			return false
		}
		// the current record is valid
		this.record++
		if this.record <= (this.response.Torow - this.response.Fromrow) {
			return true
		}
		// we reached the end of result set
		if this.response.Rows == this.response.Torow {
			// gaurd against over fill
			this.record--
			return false
		}
		// if we are here there is another batch
		this.reset()
		header, bytes, ok := this.read()
		if !ok {
			return false
		}
		// should not happen but check anyway
		// when RequestId is 0 it means we are reading published data
		if header.RequestId > 0 && header.RequestId != this.requestId {
			this.setErrorString("protocol error")
			return false
		}
		// we got another batch unmarshall the data
		this.unmarshalJSON(bytes)
	}
	return false
}

//Value returns the value within the current row for the given column name.
//If the column name does not exist, Value returns an empty string.
func (this *Client) Value(column string) string {
	ordinal, ok := this.columns[column]
	if !ok {
		return ""
	}
	return this.ValueByOrdinal(ordinal)
}

//ValueByOrdinal returns the value within the current row for the given column ordinal.
//The column ordinal represents the zero based position of the column in the Columns collection of the result set.
//If the column ordinal is out of range, ValueByOrdinal returns an empty string.
func (this *Client) ValueByOrdinal(ordinal int) string {
	if this.record < 0 || this.record >= len(this.response.Data) {
		return ""
	}
	if ordinal >= len(this.response.Data[this.record]) {
		return ""
	}
	return this.response.Data[this.record][ordinal]
}

//HasColumn determines if the column name exists in the columns collection of the result set.
func (this *Client) HasColumn(column string) bool {
	_, ok := this.columns[column]
	return ok
}

//ColumnCount returns the number of columns in the columns collection of the result set.
func (this *Client) ColumnCount() int {
	return len(this.response.Columns)
}

//Columns returns the column names in the columns collection of the result set.
func (this *Client) Columns() []string {
	return this.response.Columns
}

//WaitForPubSub waits until the pubsubsql server publishes a message for
// the subscribed Client or until the timeout interval elapses.
//Returns false when timeout interval elapses or if there is and error.
//To find out if false was returned because of an error, use Ok or Failed functions.
func (this *Client) WaitForPubSub(timeout int) bool {
	var bytes []byte
	for {
		this.reset()
		// process backlog first
		bytes = this.popBacklog()
		if len(bytes) > 0 {
			return this.unmarshalJSON(bytes)
		}
		header, temp, success, timedout := this.readTimeout(int64(timeout))
		bytes = temp
		if !success || timedout {
			return false
		}
		if header.RequestId == 0 {
			return this.unmarshalJSON(bytes)
		}
		// this is not pubsub message; are we reading abandoned cursor?
		// ignore and keep trying
	}
	return false
}

func (this *Client) popBacklog() []byte {
	element := this.backlog.Front()
	if element != nil {
		bytes := element.Value.([]byte)
		this.backlog.Remove(element)
		return bytes
	}
	return nil
}

func (this *Client) unmarshalJSON(bytes []byte) bool {
	this.rawjson = bytes
	err := json.Unmarshal(bytes, &this.response)
	if err != nil {
		this.setError(err)
		return false
	}
	if this.response.Status != "ok" {
		this.setErrorString(this.response.Msg)
		return false
	}
	this.setColumns()
	return true
}

func (this *Client) setColumns() {
	if len(this.response.Columns) == 0 {
		return
	}
	this.columns = make(map[string]int, cap(this.response.Columns))
	for ordinal, column := range this.response.Columns {
		this.columns[column] = ordinal
	}
}

func (this *Client) reset() {
	this.resetError()
	this.response.reset()
	this.rawjson = nil
	this.record = -1
}

func (this *Client) resetError() {
	this.err = ""
}

func (this *Client) setErrorString(err string) {
	this.reset()
	this.err = err
}

func (this *Client) setError(err error) {
	this.setErrorString(err.Error())
}

func (this *Client) write(message string) bool {
	this.requestId++
	if !this.rw.valid() {
		this.setErrorString("Not connected")
		return false
	}
	err := this.rw.writeHeaderAndMessage(this.requestId, []byte(message))
	if err != nil {
		this.setError(err)
		return false
	}
	return true
}

func (this *Client) readTimeout(timeout int64) (*netHeader, []byte, bool, bool) {
	if !this.rw.valid() {
		this.setErrorString("Not connected")
		return nil, nil, false, false
	}
	header, bytes, err, timedout := this.rw.readMessageTimeout(timeout)
	if timedout {
		return nil, nil, true, true
	}
	// error
	if err != nil {
		this.setError(err)
		return nil, nil, false, false
	}
	// success
	return header, bytes, true, false

}

func (this *Client) read() (*netHeader, []byte, bool) {
	var MAX_READ_TIMEOUT_MILLISECONDS int64 = 1000 * 60 * 3
	header, bytes, success, timedout := this.readTimeout(MAX_READ_TIMEOUT_MILLISECONDS)
	if timedout {
		this.setErrorString("Read timed out")
		return nil, nil, false
	}
	return header, bytes, success
}
<|endoftext|>


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

// Package docker provides a client for the Docker remote API.
//
// See http://goo.gl/mxyql for more details on the remote API.
package docker

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

const (
	apiVersion = 1.1
	userAgent  = "go-dockerclient"
)

// ErrInvalidEndpoint is the error returned by NewClient when the given
// endpoint is invalid.
var ErrInvalidEndpoint = errors.New("Invalid endpoint")

// Client is the basic type of this package. It provides methods for
// interaction with the API.
type Client struct {
	endpoint string
	client   *http.Client
}

// NewClient returns a Client instance ready for communication with the
// given server endpoint.
func NewClient(endpoint string) (*Client, error) {
	if !isValid(endpoint) {
		return nil, ErrInvalidEndpoint
	}
	return &Client{endpoint: endpoint, client: http.DefaultClient}, nil
}

func (c *Client) do(method, path string, data interface{}) ([]byte, int, error) {
	var params io.Reader
	if data != nil {
		buf, err := json.Marshal(data)
		if err != nil {
			return nil, -1, err
		}
		params = bytes.NewBuffer(buf)
	}
	req, err := http.NewRequest(method, c.getURL(path), params)
	if err != nil {
		return nil, -1, err
	}
	req.Header.Set("User-Agent", userAgent)
	if data != nil {
		req.Header.Set("Content-Type", "application/json")
	} else if method == "POST" {
		req.Header.Set("Content-Type", "plain/text")
	}
	resp, err := c.client.Do(req)
	if err != nil {
		if strings.Contains(err.Error(), "connection refused") {
			return nil, -1, fmt.Errorf("Can't connect to docker daemon. Is 'docker -d' running on this host?")
		}
		return nil, -1, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, -1, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		return nil, resp.StatusCode, newAPIClientError(resp.StatusCode, body)
	}
	return body, resp.StatusCode, nil
}

func (c *Client) stream(method, path string, in io.Reader, out io.Writer) error {
	if (method == "POST" || method == "PUT") && in == nil {
		in = bytes.NewReader([]byte{})
	}
	req, err := http.NewRequest(method, c.getURL(path), in)
	if err != nil {
		return err
	}
	req.Header.Set("User-Agent", userAgent)
	if method == "POST" {
		req.Header.Set("Content-Type", "plain/text")
	}
	resp, err := c.client.Do(req)
	if err != nil {
		if strings.Contains(err.Error(), "connection refused") {
			return fmt.Errorf("Can't connect to docker daemon. Is 'docker -d' running on this host?")
		}
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		return newAPIClientError(resp.StatusCode, body)
	}
	if resp.Header.Get("Content-Type") == "application/json" {
		dec := json.NewDecoder(resp.Body)
		for {
			var m JSONMessage
			if err := dec.Decode(&m); err == io.EOF {
				break
			} else if err != nil {
				return err
			}
			if m.Progress != "" {
				fmt.Fprintf(out, "%s %s\r", m.Status, m.Progress)
			} else if m.Error != "" {
				return fmt.Errorf(m.Error)
			} else {
				fmt.Fprintf(out, "%s\n", m.Status)
			}
		}
	} else {
		if _, err := io.Copy(out, resp.Body); err != nil {
			return err
		}
	}
	return nil
}

func (c *Client) getURL(path string) string {
	return fmt.Sprintf("%s/v%f%s", strings.TrimRight(c.endpoint, "/"), apiVersion, path)
}

type JSONMessage struct {
	Status   string `json:"status,omitempty"`
	Progress string `json:"progress,omitempty"`
	Error    string `json:"error,omitempty"`
}

func queryString(opts interface{}) string {
	if opts == nil {
		return ""
	}
	value := reflect.ValueOf(opts)
	if value.Kind() == reflect.Ptr {
		value = value.Elem()
	}
	if value.Kind() != reflect.Struct {
		return ""
	}
	items := url.Values(map[string][]string{})
	for i := 0; i < value.NumField(); i++ {
		field := value.Type().Field(i)
		if field.PkgPath != "" {
			continue
		}
		key := field.Tag.Get("qs")
		if key == "" {
			key = strings.ToLower(field.Name)
		}
		v := value.Field(i)
		switch v.Kind() {
		case reflect.Bool:
			if v.Bool() {
				items.Add(key, "1")
			}
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			if v.Int() > 0 {
				items.Add(key, strconv.FormatInt(v.Int(), 10))
			}
		case reflect.Float32, reflect.Float64:
			if v.Float() > 0 {
				items.Add(key, strconv.FormatFloat(v.Float(), 'f', -1, 64))
			}
		case reflect.String:
			if v.String() != "" {
				items.Add(key, v.String())
			}
		case reflect.Ptr:
			if !v.IsNil() {
				if b, err := json.Marshal(v.Interface()); err == nil {
					items.Add(key, string(b))
				}
			}
		}
	}
	return items.Encode()
}

type apiClientError struct {
	status  int
	message string
}

func newAPIClientError(status int, body []byte) *apiClientError {
	return &apiClientError{status: status, message: string(body)}
}

func (e *apiClientError) Error() string {
	return fmt.Sprintf("API error (%d): %s", e.status, e.message)
}

func isValid(endpoint string) bool {
	u, err := url.Parse(endpoint)
	if err != nil {
		return false
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return false
	}
	_, port, err := net.SplitHostPort(u.Host)
	if err != nil {
		if e, ok := err.(*net.AddrError); ok {
			return e.Err == "missing port in address"
		}
		return false
	}
	number, err := strconv.ParseInt(port, 10, 64)
	return err == nil && number > 0 && number < 65536
}
<commit_msg>client: some improvement in the code structure<commit_after>// Copyright 2013 Francisco Souza. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package docker provides a client for the Docker remote API.
//
// See http://goo.gl/mxyql for more details on the remote API.
package docker

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

const (
	apiVersion = 1.1
	userAgent  = "go-dockerclient"
)

var (
	// Error returned when the endpoint is not a valid HTTP URL.
	ErrInvalidEndpoint = errors.New("Invalid endpoint")

	// Error returned when the client cannot connect to the given endpoint.
	ErrConnectionRefused = errors.New("Cannot connect to Docker endpoint")
)

// Client is the basic type of this package. It provides methods for
// interaction with the API.
type Client struct {
	endpoint string
	client   *http.Client
}

// NewClient returns a Client instance ready for communication with the
// given server endpoint.
func NewClient(endpoint string) (*Client, error) {
	if !isValid(endpoint) {
		return nil, ErrInvalidEndpoint
	}
	return &Client{endpoint: endpoint, client: http.DefaultClient}, nil
}

func (c *Client) do(method, path string, data interface{}) ([]byte, int, error) {
	var params io.Reader
	if data != nil {
		buf, err := json.Marshal(data)
		if err != nil {
			return nil, -1, err
		}
		params = bytes.NewBuffer(buf)
	}
	req, err := http.NewRequest(method, c.getURL(path), params)
	if err != nil {
		return nil, -1, err
	}
	req.Header.Set("User-Agent", userAgent)
	if data != nil {
		req.Header.Set("Content-Type", "application/json")
	} else if method == "POST" {
		req.Header.Set("Content-Type", "plain/text")
	}
	resp, err := c.client.Do(req)
	if err != nil {
		if strings.Contains(err.Error(), "connection refused") {
			return nil, -1, ErrConnectionRefused
		}
		return nil, -1, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, -1, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		return nil, resp.StatusCode, newAPIClientError(resp.StatusCode, body)
	}
	return body, resp.StatusCode, nil
}

func (c *Client) stream(method, path string, in io.Reader, out io.Writer) error {
	if (method == "POST" || method == "PUT") && in == nil {
		in = bytes.NewReader(nil)
	}
	req, err := http.NewRequest(method, c.getURL(path), in)
	if err != nil {
		return err
	}
	req.Header.Set("User-Agent", userAgent)
	if method == "POST" {
		req.Header.Set("Content-Type", "plain/text")
	}
	resp, err := c.client.Do(req)
	if err != nil {
		if strings.Contains(err.Error(), "connection refused") {
			return ErrConnectionRefused
		}
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		return newAPIClientError(resp.StatusCode, body)
	}
	if resp.Header.Get("Content-Type") == "application/json" {
		dec := json.NewDecoder(resp.Body)
		for {
			var m JSONMessage
			if err := dec.Decode(&m); err == io.EOF {
				break
			} else if err != nil {
				return err
			}
			if m.Progress != "" {
				fmt.Fprintf(out, "%s %s\r", m.Status, m.Progress)
			} else if m.Error != "" {
				return errors.New(m.Error)
			} else {
				fmt.Fprintf(out, "%s\n", m.Status)
			}
		}
	} else {
		if _, err := io.Copy(out, resp.Body); err != nil {
			return err
		}
	}
	return nil
}

func (c *Client) getURL(path string) string {
	return fmt.Sprintf("%s/v%f%s", strings.TrimRight(c.endpoint, "/"), apiVersion, path)
}

type JSONMessage struct {
	Status   string `json:"status,omitempty"`
	Progress string `json:"progress,omitempty"`
	Error    string `json:"error,omitempty"`
}

func queryString(opts interface{}) string {
	if opts == nil {
		return ""
	}
	value := reflect.ValueOf(opts)
	if value.Kind() == reflect.Ptr {
		value = value.Elem()
	}
	if value.Kind() != reflect.Struct {
		return ""
	}
	items := url.Values(map[string][]string{})
	for i := 0; i < value.NumField(); i++ {
		field := value.Type().Field(i)
		if field.PkgPath != "" {
			continue
		}
		key := field.Tag.Get("qs")
		if key == "" {
			key = strings.ToLower(field.Name)
		}
		v := value.Field(i)
		switch v.Kind() {
		case reflect.Bool:
			if v.Bool() {
				items.Add(key, "1")
			}
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			if v.Int() > 0 {
				items.Add(key, strconv.FormatInt(v.Int(), 10))
			}
		case reflect.Float32, reflect.Float64:
			if v.Float() > 0 {
				items.Add(key, strconv.FormatFloat(v.Float(), 'f', -1, 64))
			}
		case reflect.String:
			if v.String() != "" {
				items.Add(key, v.String())
			}
		case reflect.Ptr:
			if !v.IsNil() {
				if b, err := json.Marshal(v.Interface()); err == nil {
					items.Add(key, string(b))
				}
			}
		}
	}
	return items.Encode()
}

type apiClientError struct {
	status  int
	message string
}

func newAPIClientError(status int, body []byte) *apiClientError {
	return &apiClientError{status: status, message: string(body)}
}

func (e *apiClientError) Error() string {
	return fmt.Sprintf("API error (%d): %s", e.status, e.message)
}

func isValid(endpoint string) bool {
	u, err := url.Parse(endpoint)
	if err != nil {
		return false
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return false
	}
	_, port, err := net.SplitHostPort(u.Host)
	if err != nil {
		if e, ok := err.(*net.AddrError); ok {
			return e.Err == "missing port in address"
		}
		return false
	}
	number, err := strconv.ParseInt(port, 10, 64)
	return err == nil && number > 0 && number < 65536
}
<|endoftext|>


<commit_before>package gosseract

// #if __FreeBSD__ >= 10
// #cgo LDFLAGS: -L/usr/local/lib -llept -ltesseract
// #else
// #cgo LDFLAGS: -llept -ltesseract
// #endif
// #include <stdlib.h>
// #include "tessbridge.h"
import "C"
import (
	"strings"
	"unsafe"
)

// Version returns the version of Tesseract-OCR
func Version() string {
	api := C.Create()
	defer C.Free(api)
	version := C.Version(api)
	return C.GoString(version)
}

// Client is argument builder for tesseract::TessBaseAPI.
type Client struct {
	api C.TessBaseAPI

	// Trim specifies characters to trim, which would be trimed from result string.
	// As results of OCR, text often contains unnecessary characters, such as newlines, on the head/foot of string.
	// If `Trim` is set, this client will remove specified characters from the result.
	Trim bool

	// TessdataPrefix can indicate directory path to `tessdata`.
	// It is set `/usr/local/share/tessdata/` or something like that, as default.
	// TODO: Implement and test
	TessdataPrefix *string

	// Languages are languages to be detected. If not specified, it's gonna be "eng".
	Languages []string

	// ImagePath is just path to image file to be processed OCR.
	ImagePath string

	// Variables is just a pool to evaluate "tesseract::TessBaseAPI->SetVariable" in delay.
	// TODO: Think if it should be public, or private property.
	Variables map[string]string

	// PageSegMode is a mode for page layout analysis.
	// See https://github.com/otiai10/gosseract/issues/52 for more information.
	PageSegMode *PageSegMode
}

// NewClient construct new Client. It's due to caller to Close this client.
func NewClient() *Client {
	client := &Client{
		api:       C.Create(),
		Variables: map[string]string{},
		Trim:      true,
	}
	return client
}

// Close frees allocated API. This MUST be called for ANY client constructed by "NewClient" function.
func (c *Client) Close() (err error) {
	// defer func() {
	// 	if e := recover(); e != nil {
	// 		err = fmt.Errorf("%v", e)
	// 	}
	// }()
	C.Free(c.api)
	return err
}

// SetImage sets path to image file to be processed OCR.
func (c *Client) SetImage(imagepath string) *Client {
	c.ImagePath = imagepath
	return c
}

// SetLanguage sets languages to use. English as default.
func (c *Client) SetLanguage(langs ...string) *Client {
	c.Languages = langs
	return c
}

// SetWhitelist sets whitelist chars.
// See official documentation for whitelist here https://github.com/tesseract-ocr/tesseract/wiki/ImproveQuality#dictionaries-word-lists-and-patterns
func (c *Client) SetWhitelist(whitelist string) *Client {
	return c.SetVariable("tessedit_char_whitelist", whitelist)
}

// SetVariable sets parameters, representing tesseract::TessBaseAPI->SetVariable.
// See official documentation here https://zdenop.github.io/tesseract-doc/classtesseract_1_1_tess_base_a_p_i.html#a2e09259c558c6d8e0f7e523cbaf5adf5
func (c *Client) SetVariable(key, value string) *Client {
	c.Variables[key] = value
	return c
}

// SetPageSegMode sets "Page Segmentation Mode" (PSM) to detect layout of characters.
// See official documentation for PSM here https://github.com/tesseract-ocr/tesseract/wiki/ImproveQuality#page-segmentation-method
func (c *Client) SetPageSegMode(mode PageSegMode) *Client {
	c.PageSegMode = &mode
	return c
}

// Initialize tesseract::TessBaseAPI
// TODO: add tessdata prefix
func (c *Client) init() {
	if len(c.Languages) == 0 {
		C.Init(c.api, nil, nil)
	} else {
		langs := C.CString(strings.Join(c.Languages, "+"))
		defer C.free(unsafe.Pointer(langs))
		C.Init(c.api, nil, langs)
	}
}

// Text finally initialize tesseract::TessBaseAPI, execute OCR and extract text detected as string.
func (c *Client) Text() (string, error) {

	// Defer recover and make error
	var err error
	// TODO: Handle and recover errors by Cgo.
	// defer func() {
	// 	if e := recover(); e != nil {
	// 		err = fmt.Errorf("%v", e)
	// 	}
	// }()

	c.init()

	// Set Image by giving path
	imagepath := C.CString(c.ImagePath)
	defer C.free(unsafe.Pointer(imagepath))
	C.SetImage(c.api, imagepath)

	for key, value := range c.Variables {
		k, v := C.CString(key), C.CString(value)
		defer C.free(unsafe.Pointer(k))
		defer C.free(unsafe.Pointer(v))
		C.SetVariable(c.api, k, v)
	}

	if c.PageSegMode != nil {
		mode := C.int(*c.PageSegMode)
		C.SetPageSegMode(c.api, mode)
	}

	// Get text by execuitng
	out := C.GoString(C.UTF8Text(c.api))

	// Trim result if needed
	if c.Trim {
		out = strings.Trim(out, "\n")
	}

	return out, err
}
<commit_msg>add CXXFLAGS for compiling with ubuntu<commit_after>package gosseract

// #cgo CXXFLAGS: -std=c++11
// #if __FreeBSD__ >= 10
// #cgo LDFLAGS: -L/usr/local/lib -llept -ltesseract
// #else
// #cgo LDFLAGS: -llept -ltesseract
// #endif
// #include <stdlib.h>
// #include "tessbridge.h"
import "C"
import (
	"strings"
	"unsafe"
)

// Version returns the version of Tesseract-OCR
func Version() string {
	api := C.Create()
	defer C.Free(api)
	version := C.Version(api)
	return C.GoString(version)
}

// Client is argument builder for tesseract::TessBaseAPI.
type Client struct {
	api C.TessBaseAPI

	// Trim specifies characters to trim, which would be trimed from result string.
	// As results of OCR, text often contains unnecessary characters, such as newlines, on the head/foot of string.
	// If `Trim` is set, this client will remove specified characters from the result.
	Trim bool

	// TessdataPrefix can indicate directory path to `tessdata`.
	// It is set `/usr/local/share/tessdata/` or something like that, as default.
	// TODO: Implement and test
	TessdataPrefix *string

	// Languages are languages to be detected. If not specified, it's gonna be "eng".
	Languages []string

	// ImagePath is just path to image file to be processed OCR.
	ImagePath string

	// Variables is just a pool to evaluate "tesseract::TessBaseAPI->SetVariable" in delay.
	// TODO: Think if it should be public, or private property.
	Variables map[string]string

	// PageSegMode is a mode for page layout analysis.
	// See https://github.com/otiai10/gosseract/issues/52 for more information.
	PageSegMode *PageSegMode
}

// NewClient construct new Client. It's due to caller to Close this client.
func NewClient() *Client {
	client := &Client{
		api:       C.Create(),
		Variables: map[string]string{},
		Trim:      true,
	}
	return client
}

// Close frees allocated API. This MUST be called for ANY client constructed by "NewClient" function.
func (c *Client) Close() (err error) {
	// defer func() {
	// 	if e := recover(); e != nil {
	// 		err = fmt.Errorf("%v", e)
	// 	}
	// }()
	C.Free(c.api)
	return err
}

// SetImage sets path to image file to be processed OCR.
func (c *Client) SetImage(imagepath string) *Client {
	c.ImagePath = imagepath
	return c
}

// SetLanguage sets languages to use. English as default.
func (c *Client) SetLanguage(langs ...string) *Client {
	c.Languages = langs
	return c
}

// SetWhitelist sets whitelist chars.
// See official documentation for whitelist here https://github.com/tesseract-ocr/tesseract/wiki/ImproveQuality#dictionaries-word-lists-and-patterns
func (c *Client) SetWhitelist(whitelist string) *Client {
	return c.SetVariable("tessedit_char_whitelist", whitelist)
}

// SetVariable sets parameters, representing tesseract::TessBaseAPI->SetVariable.
// See official documentation here https://zdenop.github.io/tesseract-doc/classtesseract_1_1_tess_base_a_p_i.html#a2e09259c558c6d8e0f7e523cbaf5adf5
func (c *Client) SetVariable(key, value string) *Client {
	c.Variables[key] = value
	return c
}

// SetPageSegMode sets "Page Segmentation Mode" (PSM) to detect layout of characters.
// See official documentation for PSM here https://github.com/tesseract-ocr/tesseract/wiki/ImproveQuality#page-segmentation-method
func (c *Client) SetPageSegMode(mode PageSegMode) *Client {
	c.PageSegMode = &mode
	return c
}

// Initialize tesseract::TessBaseAPI
// TODO: add tessdata prefix
func (c *Client) init() {
	if len(c.Languages) == 0 {
		C.Init(c.api, nil, nil)
	} else {
		langs := C.CString(strings.Join(c.Languages, "+"))
		defer C.free(unsafe.Pointer(langs))
		C.Init(c.api, nil, langs)
	}
}

// Text finally initialize tesseract::TessBaseAPI, execute OCR and extract text detected as string.
func (c *Client) Text() (string, error) {

	// Defer recover and make error
	var err error
	// TODO: Handle and recover errors by Cgo.
	// defer func() {
	// 	if e := recover(); e != nil {
	// 		err = fmt.Errorf("%v", e)
	// 	}
	// }()

	c.init()

	// Set Image by giving path
	imagepath := C.CString(c.ImagePath)
	defer C.free(unsafe.Pointer(imagepath))
	C.SetImage(c.api, imagepath)

	for key, value := range c.Variables {
		k, v := C.CString(key), C.CString(value)
		defer C.free(unsafe.Pointer(k))
		defer C.free(unsafe.Pointer(v))
		C.SetVariable(c.api, k, v)
	}

	if c.PageSegMode != nil {
		mode := C.int(*c.PageSegMode)
		C.SetPageSegMode(c.api, mode)
	}

	// Get text by execuitng
	out := C.GoString(C.UTF8Text(c.api))

	// Trim result if needed
	if c.Trim {
		out = strings.Trim(out, "\n")
	}

	return out, err
}
<|endoftext|>


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

package keywhizfs

import (
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"

	klog "github.com/square/keywhiz-fs/log"
)

// clientRefresh is the rate the client reloads itself in the background.
const clientRefresh = 10 * time.Minute

// Cipher suites enabled in the client. No RC4 or 3DES.
var ciphers = []uint16{
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
	tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
	tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
	tls.TLS_RSA_WITH_AES_128_CBC_SHA,
	tls.TLS_RSA_WITH_AES_256_CBC_SHA,
}

// Client basic struct.
type Client struct {
	*klog.Logger
	http func() *http.Client
	url  *url.URL
}

// httpClientParams are values necessary for constructing a TLS client.
type httpClientParams struct {
	certFile,
	keyFile,
	caFile string
	timeout time.Duration
}

// NewClient produces a read-to-use client struct given PEM-encoded certificate file, key file, and
// ca file with the list of trusted certificate authorities.
func NewClient(certFile, keyFile, caFile string, serverURL *url.URL, timeout time.Duration, logConfig klog.Config, ping bool) (client Client) {
	logger := klog.New("kwfs_client", logConfig)
	params := httpClientParams{certFile, keyFile, caFile, timeout}

	reqc := make(chan http.Client)

	// Getter from channel.
	getClient := func() *http.Client {
		client := <-reqc
		return &(client)
	}

	initial, err := params.buildClient()
	if err != nil {
		panic(err)
	}

	// Asynchronously updates client and owns current reference.
	go func() {
		current := *initial
		ticker := time.Tick(clientRefresh)
		for {
			select {
			case t := <-ticker: // Periodically update client.
				logger.Infof("Updating http client at %v", t)
				if c, err := params.buildClient(); err != nil {
					logger.Errorf("Error refreshing http client: %v", err)
				} else {
					current = *c
				}
			case reqc <- current: // Service request for current client.
			}
		}
	}()

	client = Client{logger, getClient, serverURL}
	if ping {
		if _, ok := client.SecretList(); !ok {
			log.Fatalf("Failed startup /secrets ping to %v", client.url)
		}
	}

	return client
}

// RawSecret returns raw JSON from requesting a secret.
func (c Client) RawSecret(name string) (data []byte, ok bool) {
	now := time.Now()
	// note: path.Join does not know how to properly escape for URLs!
	t := *c.url
	t.Path = path.Join(c.url.Path, "secret", name)
	resp, err := c.http().Get(t.String())
	if err != nil {
		c.Errorf("Error retrieving secret %v: %v", name, err)
		return nil, false
	}
	c.Infof("GET /secret/%v %d %v", name, resp.StatusCode, time.Since(now))
	defer resp.Body.Close()

	data, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		c.Errorf("Error reading response body for secret %v: %v", name, err)
		return nil, false
	}

	switch resp.StatusCode {
	case 200:
		return data, true
	case 404:
		c.Warnf("Secret %v not found", name)
		return nil, false
	default:
		msg := strings.Join(strings.Split(string(data), "\n"), " ")
		c.Errorf("Bad response code getting secret %v: (status=%v, msg='%s')", name, resp.StatusCode, msg)
		return nil, false
	}
}

// Secret returns an unmarshalled Secret struct after requesting a secret.
func (c Client) Secret(name string) (secret *Secret, ok bool) {
	data, ok := c.RawSecret(name)
	if !ok {
		return nil, false
	}

	secret, err := ParseSecret(data)
	if err != nil {
		c.Errorf("Error decoding retrieved secret %v: %v", name, err)
		return nil, false
	}

	return secret, true
}

// RawSecretList returns raw JSON from requesting a listing of secrets.
func (c Client) RawSecretList() (data []byte, ok bool) {
	now := time.Now()
	t := *c.url
	t.Path = path.Join(c.url.Path, "secrets")
	resp, err := c.http().Get(t.String())
	if err != nil {
		c.Errorf("Error retrieving secrets: %v", err)
		return nil, false
	}
	c.Infof("GET /secrets %d %v", resp.StatusCode, time.Since(now))
	defer resp.Body.Close()

	data, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		c.Errorf("Error reading response body for secrets: %v", err)
		return nil, false
	}

	if resp.StatusCode != 200 {
		msg := strings.Join(strings.Split(string(data), "\n"), " ")
		c.Errorf("Bad response code getting secrets: (status=%v, msg='%s')", resp.StatusCode, msg)
		return nil, false
	}
	return data, true
}

// SecretList returns a slice of unmarshalled Secret structs after requesting a listing of secrets.
func (c Client) SecretList() (secrets []Secret, ok bool) {
	data, ok := c.RawSecretList()
	if !ok {
		return nil, false
	}

	secrets, err := ParseSecretList(data)
	if err != nil {
		c.Errorf("Error decoding retrieved secrets: %v", err)
		return nil, false
	}
	return secrets, true
}

// buildClient constructs a new TLS client.
func (p httpClientParams) buildClient() (client *http.Client, err error) {
	keyPair, err := tls.LoadX509KeyPair(p.certFile, p.keyFile)
	if err != nil {
		return
	}

	caCert, err := ioutil.ReadFile(p.caFile)
	if err != nil {
		return
	}
	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caCert)

	config := &tls.Config{
		Certificates: []tls.Certificate{keyPair},
		RootCAs:      caCertPool,
		MinVersion:   tls.VersionTLS12, // TLSv1.2 and up is required
		CipherSuites: ciphers,
	}
	config.BuildNameToCertificate()
	transport := &http.Transport{TLSClientConfig: config}
	return &http.Client{Transport: transport, Timeout: p.timeout}, nil
}
<commit_msg>Avoid copying http.Client around, just pass reference<commit_after>// Copyright 2015 Square Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package keywhizfs

import (
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"path"
	"strings"
	"time"

	klog "github.com/square/keywhiz-fs/log"
)

// clientRefresh is the rate the client reloads itself in the background.
const clientRefresh = 10 * time.Minute

// Cipher suites enabled in the client. No RC4 or 3DES.
var ciphers = []uint16{
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
	tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
	tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
	tls.TLS_RSA_WITH_AES_128_CBC_SHA,
	tls.TLS_RSA_WITH_AES_256_CBC_SHA,
}

// Client basic struct.
type Client struct {
	*klog.Logger
	http func() *http.Client
	url  *url.URL
}

// httpClientParams are values necessary for constructing a TLS client.
type httpClientParams struct {
	certFile,
	keyFile,
	caFile string
	timeout time.Duration
}

// NewClient produces a read-to-use client struct given PEM-encoded certificate file, key file, and
// ca file with the list of trusted certificate authorities.
func NewClient(certFile, keyFile, caFile string, serverURL *url.URL, timeout time.Duration, logConfig klog.Config, ping bool) (client Client) {
	logger := klog.New("kwfs_client", logConfig)
	params := httpClientParams{certFile, keyFile, caFile, timeout}

	reqc := make(chan *http.Client)

	// Getter from channel.
	getClient := func() *http.Client {
		client := <-reqc
		return client
	}

	initial, err := params.buildClient()
	if err != nil {
		panic(err)
	}

	// Asynchronously updates client and owns current reference.
	go func() {
		current := initial
		ticker := time.Tick(clientRefresh)
		for {
			select {
			case t := <-ticker: // Periodically update client.
				logger.Infof("Updating http client at %v", t)
				if client, err := params.buildClient(); err != nil {
					logger.Errorf("Error refreshing http client: %v", err)
				} else {
					current = client
				}
			case reqc <- current: // Service request for current client.
			}
		}
	}()

	client = Client{logger, getClient, serverURL}
	if ping {
		if _, ok := client.SecretList(); !ok {
			log.Fatalf("Failed startup /secrets ping to %v", client.url)
		}
	}

	return client
}

// RawSecret returns raw JSON from requesting a secret.
func (c Client) RawSecret(name string) (data []byte, ok bool) {
	now := time.Now()
	// note: path.Join does not know how to properly escape for URLs!
	t := *c.url
	t.Path = path.Join(c.url.Path, "secret", name)
	resp, err := c.http().Get(t.String())
	if err != nil {
		c.Errorf("Error retrieving secret %v: %v", name, err)
		return nil, false
	}
	c.Infof("GET /secret/%v %d %v", name, resp.StatusCode, time.Since(now))
	defer resp.Body.Close()

	data, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		c.Errorf("Error reading response body for secret %v: %v", name, err)
		return nil, false
	}

	switch resp.StatusCode {
	case 200:
		return data, true
	case 404:
		c.Warnf("Secret %v not found", name)
		return nil, false
	default:
		msg := strings.Join(strings.Split(string(data), "\n"), " ")
		c.Errorf("Bad response code getting secret %v: (status=%v, msg='%s')", name, resp.StatusCode, msg)
		return nil, false
	}
}

// Secret returns an unmarshalled Secret struct after requesting a secret.
func (c Client) Secret(name string) (secret *Secret, ok bool) {
	data, ok := c.RawSecret(name)
	if !ok {
		return nil, false
	}

	secret, err := ParseSecret(data)
	if err != nil {
		c.Errorf("Error decoding retrieved secret %v: %v", name, err)
		return nil, false
	}

	return secret, true
}

// RawSecretList returns raw JSON from requesting a listing of secrets.
func (c Client) RawSecretList() (data []byte, ok bool) {
	now := time.Now()
	t := *c.url
	t.Path = path.Join(c.url.Path, "secrets")
	resp, err := c.http().Get(t.String())
	if err != nil {
		c.Errorf("Error retrieving secrets: %v", err)
		return nil, false
	}
	c.Infof("GET /secrets %d %v", resp.StatusCode, time.Since(now))
	defer resp.Body.Close()

	data, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		c.Errorf("Error reading response body for secrets: %v", err)
		return nil, false
	}

	if resp.StatusCode != 200 {
		msg := strings.Join(strings.Split(string(data), "\n"), " ")
		c.Errorf("Bad response code getting secrets: (status=%v, msg='%s')", resp.StatusCode, msg)
		return nil, false
	}
	return data, true
}

// SecretList returns a slice of unmarshalled Secret structs after requesting a listing of secrets.
func (c Client) SecretList() (secrets []Secret, ok bool) {
	data, ok := c.RawSecretList()
	if !ok {
		return nil, false
	}

	secrets, err := ParseSecretList(data)
	if err != nil {
		c.Errorf("Error decoding retrieved secrets: %v", err)
		return nil, false
	}
	return secrets, true
}

// buildClient constructs a new TLS client.
func (p httpClientParams) buildClient() (client *http.Client, err error) {
	keyPair, err := tls.LoadX509KeyPair(p.certFile, p.keyFile)
	if err != nil {
		return
	}

	caCert, err := ioutil.ReadFile(p.caFile)
	if err != nil {
		return
	}
	caCertPool := x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caCert)

	config := &tls.Config{
		Certificates: []tls.Certificate{keyPair},
		RootCAs:      caCertPool,
		MinVersion:   tls.VersionTLS12, // TLSv1.2 and up is required
		CipherSuites: ciphers,
	}
	config.BuildNameToCertificate()
	transport := &http.Transport{TLSClientConfig: config}
	return &http.Client{Transport: transport, Timeout: p.timeout}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"strings"
)

type Client struct {
	incomingConn net.Conn
	pass         string
	username     string
}

func newClient(conn net.Conn) Client {
	return Client{
		incomingConn: conn,
	}
}

func (c *Client) handleMessage(line string) {
	log.Debug(line)
	if strings.HasPrefix(line, "PASS ") {
        c.pass = strings.Split(line, "PASS ")[1]
	}
}
<commit_msg>use switch for handleMessage<commit_after>package main

import (
	"net"
	"strings"
)

type Client struct {
	incomingConn net.Conn
	pass         string
	username     string
}

func newClient(conn net.Conn) Client {
	return Client{
		incomingConn: conn,
	}
}

func (c *Client) handleMessage(line string) {
	log.Debug(line)
	spl := strings.Split(line, " ")
	command := spl[0]
	switch command {
	case "PASS":
		c.pass = spl[1]
	}
}
<|endoftext|>


<commit_before>package dns

// A client implementation.

import (
	"bytes"
	"io"
	"net"
	"time"
)

const dnsTimeout time.Duration = 2 * time.Second
const tcpIdleTimeout time.Duration = 8 * time.Second

// A Conn represents a connection to a DNS server.
type Conn struct {
	net.Conn                         // a net.Conn holding the connection
	UDPSize        uint16            // minimum receive buffer for UDP messages
	TsigSecret     map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	rtt            time.Duration
	t              time.Time
	tsigRequestMAC string
}

// A Client defines parameters for a DNS client.
type Client struct {
	Net            string            // if "tcp" a TCP query will be initiated, otherwise an UDP one (default is "" for UDP)
	UDPSize        uint16            // minimum receive buffer for UDP messages
	DialTimeout    time.Duration     // net.DialTimeout, defaults to 2 seconds
	ReadTimeout    time.Duration     // net.Conn.SetReadTimeout value for connections, defaults to 2 seconds
	WriteTimeout   time.Duration     // net.Conn.SetWriteTimeout value for connections, defaults to 2 seconds
	TsigSecret     map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	SingleInflight bool              // if true suppress multiple outstanding queries for the same Qname, Qtype and Qclass
	group          singleflight
}

// Exchange performs a synchronous UDP query. It sends the message m to the address
// contained in a and waits for an reply. Exchange does not retry a failed query, nor
// will it fall back to TCP in case of truncation.
// If you need to send a DNS message on an already existing connection, you can use the
// following:
//
//	co := &dns.Conn{Conn: c} // c is your net.Conn
//	co.WriteMsg(m)
//	in, err  := co.ReadMsg()
//	co.Close()
//
func Exchange(m *Msg, a string) (r *Msg, err error) {
	var co *Conn
	co, err = DialTimeout("udp", a, dnsTimeout)
	if err != nil {
		return nil, err
	}

	defer co.Close()
	co.SetReadDeadline(time.Now().Add(dnsTimeout))
	co.SetWriteDeadline(time.Now().Add(dnsTimeout))
	if err = co.WriteMsg(m); err != nil {
		return nil, err
	}
	r, err = co.ReadMsg()
	return r, err
}

// ExchangeConn performs a synchronous query. It sends the message m via the connection
// c and waits for a reply. The connection c is not closed by ExchangeConn.
// This function is going away, but can easily be mimicked:
//
//	co := &dns.Conn{Conn: c} // c is your net.Conn
//	co.WriteMsg(m)
//	in, _  := co.ReadMsg()
//	co.Close()
//
func ExchangeConn(c net.Conn, m *Msg) (r *Msg, err error) {
	println("dns: this function is deprecated")
	co := new(Conn)
	co.Conn = c
	if err = co.WriteMsg(m); err != nil {
		return nil, err
	}
	r, err = co.ReadMsg()
	return r, err
}

// Exchange performs an synchronous query. It sends the message m to the address
// contained in a and waits for an reply. Basic use pattern with a *dns.Client:
//
//	c := new(dns.Client)
//	in, rtt, err := c.Exchange(message, "127.0.0.1:53")
//
// Exchange does not retry a failed query, nor will it fall back to TCP in
// case of truncation.
func (c *Client) Exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	if !c.SingleInflight {
		return c.exchange(m, a)
	}
	// This adds a bunch of garbage, TODO(miek).
	t := "nop"
	if t1, ok := TypeToString[m.Question[0].Qtype]; ok {
		t = t1
	}
	cl := "nop"
	if cl1, ok := ClassToString[m.Question[0].Qclass]; ok {
		cl = cl1
	}
	r, rtt, err, shared := c.group.Do(m.Question[0].Name+t+cl, func() (*Msg, time.Duration, error) {
		return c.exchange(m, a)
	})
	if err != nil {
		return r, rtt, err
	}
	if shared {
		return r.Copy(), rtt, nil
	}
	return r, rtt, nil
}

func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	timeout := dnsTimeout
	var co *Conn
	if c.DialTimeout != 0 {
		timeout = c.DialTimeout
	}
	if c.Net == "" {
		co, err = DialTimeout("udp", a, timeout)
	} else {
		co, err = DialTimeout(c.Net, a, timeout)
	}
	if err != nil {
		return nil, 0, err
	}
	timeout = dnsTimeout
	if c.ReadTimeout != 0 {
		timeout = c.ReadTimeout
	}
	co.SetReadDeadline(time.Now().Add(timeout))
	timeout = dnsTimeout
	if c.WriteTimeout != 0 {
		timeout = c.WriteTimeout
	}
	co.SetWriteDeadline(time.Now().Add(timeout))
	defer co.Close()
	opt := m.IsEdns0()
	// If EDNS0 is used use that for size.
	if opt != nil && opt.UDPSize() >= MinMsgSize {
		co.UDPSize = opt.UDPSize()
	}
	// Otherwise use the client's configured UDP size.
	if opt == nil && c.UDPSize >= MinMsgSize {
		co.UDPSize = c.UDPSize
	}
	co.TsigSecret = c.TsigSecret
	if err = co.WriteMsg(m); err != nil {
		return nil, 0, err
	}
	r, err = co.ReadMsg()
	return r, co.rtt, err
}

// ReadMsg reads a message from the connection co.
// If the received message contains a TSIG record the transaction
// signature is verified.
func (co *Conn) ReadMsg() (*Msg, error) {
	var p []byte
	m := new(Msg)
	if _, ok := co.Conn.(*net.TCPConn); ok {
		p = make([]byte, MaxMsgSize)
	} else {
		if co.UDPSize >= 512 {
			p = make([]byte, co.UDPSize)
		} else {
			p = make([]byte, MinMsgSize)
		}
	}
	n, err := co.Read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	co.rtt = time.Since(co.t)
	if t := m.IsTsig(); t != nil {
		if _, ok := co.TsigSecret[t.Hdr.Name]; !ok {
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		err = TsigVerify(p, co.TsigSecret[t.Hdr.Name], co.tsigRequestMAC, false)
	}
	return m, err
}

// Read implements the net.Conn read method.
func (co *Conn) Read(p []byte) (n int, err error) {
	if co.Conn == nil {
		return 0, ErrConnEmpty
	}
	if len(p) < 2 {
		return 0, io.ErrShortBuffer
	}
	if t, ok := co.Conn.(*net.TCPConn); ok {
		n, err = t.Read(p[0:2])
		if err != nil || n != 2 {
			return n, err
		}
		l, _ := unpackUint16(p[0:2], 0)
		if l == 0 {
			return 0, ErrShortRead
		}
		if int(l) > len(p) {
			return int(l), io.ErrShortBuffer
		}
		n, err = t.Read(p[:l])
		if err != nil {
			return n, err
		}
		i := n
		for i < int(l) {
			j, err := t.Read(p[i:int(l)])
			if err != nil {
				return i, err
			}
			i += j
		}
		n = i
		return n, err
	}
	// UDP connection
	n, err = co.Conn.Read(p)
	if err != nil {
		return n, err
	}
	return n, err
}

// WriteMsg sends a message throught the connection co.
// If the message m contains a TSIG record the transaction
// signature is calculated.
func (co *Conn) WriteMsg(m *Msg) (err error) {
	var out []byte
	if t := m.IsTsig(); t != nil {
		mac := ""
		if _, ok := co.TsigSecret[t.Hdr.Name]; !ok {
			return ErrSecret
		}
		out, mac, err = TsigGenerate(m, co.TsigSecret[t.Hdr.Name], co.tsigRequestMAC, false)
		// Set for the next read, allthough only used in zone transfers
		co.tsigRequestMAC = mac
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	co.t = time.Now()
	if _, err = co.Write(out); err != nil {
		return err
	}
	return nil
}

// Write implements the net.Conn Write method.
func (co *Conn) Write(p []byte) (n int, err error) {
	if t, ok := co.Conn.(*net.TCPConn); ok {
		lp := len(p)
		if lp < 2 {
			return 0, io.ErrShortBuffer
		}
		if lp > MaxMsgSize {
			return 0, &Error{err: "message too large"}
		}
		l := make([]byte, 2, lp+2)
		l[0], l[1] = packUint16(uint16(lp))
		p = append(l, p...)
		n, err := io.Copy(t, bytes.NewReader(p))
		return int(n), err
	}
	n, err = co.Conn.(*net.UDPConn).Write(p)
	return n, err
}

// Dial connects to the address on the named network.
func Dial(network, address string) (conn *Conn, err error) {
	conn = new(Conn)
	conn.Conn, err = net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return conn, nil
}

// DialTimeout acts like Dial but takes a timeout.
func DialTimeout(network, address string, timeout time.Duration) (conn *Conn, err error) {
	conn = new(Conn)
	conn.Conn, err = net.DialTimeout(network, address, timeout)
	if err != nil {
		return nil, err
	}
	return conn, nil
}

// Close implements the net.Conn Close method.
func (co *Conn) Close() error { return co.Conn.Close() }

// LocalAddr implements the net.Conn LocalAddr method.
func (co *Conn) LocalAddr() net.Addr { return co.Conn.LocalAddr() }

// RemoteAddr implements the net.Conn RemoteAddr method.
func (co *Conn) RemoteAddr() net.Addr { return co.Conn.RemoteAddr() }

// SetDeadline implements the net.Conn SetDeadline method.
func (co *Conn) SetDeadline(t time.Time) error { return co.Conn.SetDeadline(t) }

// SetReadDeadline implements the net.Conn SetReadDeadline method.
func (co *Conn) SetReadDeadline(t time.Time) error { return co.Conn.SetReadDeadline(t) }

// SetWriteDeadline implements the net.Conn SetWriteDeadline method.
func (co *Conn) SetWriteDeadline(t time.Time) error { return co.Conn.SetWriteDeadline(t) }
<commit_msg>Check EDNS0 bufsize in Exchange()<commit_after>package dns

// A client implementation.

import (
	"bytes"
	"io"
	"net"
	"time"
)

const dnsTimeout time.Duration = 2 * time.Second
const tcpIdleTimeout time.Duration = 8 * time.Second

// A Conn represents a connection to a DNS server.
type Conn struct {
	net.Conn                         // a net.Conn holding the connection
	UDPSize        uint16            // minimum receive buffer for UDP messages
	TsigSecret     map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	rtt            time.Duration
	t              time.Time
	tsigRequestMAC string
}

// A Client defines parameters for a DNS client.
type Client struct {
	Net            string            // if "tcp" a TCP query will be initiated, otherwise an UDP one (default is "" for UDP)
	UDPSize        uint16            // minimum receive buffer for UDP messages
	DialTimeout    time.Duration     // net.DialTimeout, defaults to 2 seconds
	ReadTimeout    time.Duration     // net.Conn.SetReadTimeout value for connections, defaults to 2 seconds
	WriteTimeout   time.Duration     // net.Conn.SetWriteTimeout value for connections, defaults to 2 seconds
	TsigSecret     map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	SingleInflight bool              // if true suppress multiple outstanding queries for the same Qname, Qtype and Qclass
	group          singleflight
}

// Exchange performs a synchronous UDP query. It sends the message m to the address
// contained in a and waits for an reply. Exchange does not retry a failed query, nor
// will it fall back to TCP in case of truncation.
// If you need to send a DNS message on an already existing connection, you can use the
// following:
//
//	co := &dns.Conn{Conn: c} // c is your net.Conn
//	co.WriteMsg(m)
//	in, err  := co.ReadMsg()
//	co.Close()
//
func Exchange(m *Msg, a string) (r *Msg, err error) {
	var co *Conn
	co, err = DialTimeout("udp", a, dnsTimeout)
	if err != nil {
		return nil, err
	}

	defer co.Close()
	co.SetReadDeadline(time.Now().Add(dnsTimeout))
	co.SetWriteDeadline(time.Now().Add(dnsTimeout))

	opt := m.IsEdns0()
	// If EDNS0 is used use that for size.
	if opt != nil && opt.UDPSize() >= MinMsgSize {
		co.UDPSize = opt.UDPSize()
	}

	if err = co.WriteMsg(m); err != nil {
		return nil, err
	}
	r, err = co.ReadMsg()
	return r, err
}

// ExchangeConn performs a synchronous query. It sends the message m via the connection
// c and waits for a reply. The connection c is not closed by ExchangeConn.
// This function is going away, but can easily be mimicked:
//
//	co := &dns.Conn{Conn: c} // c is your net.Conn
//	co.WriteMsg(m)
//	in, _  := co.ReadMsg()
//	co.Close()
//
func ExchangeConn(c net.Conn, m *Msg) (r *Msg, err error) {
	println("dns: this function is deprecated")
	co := new(Conn)
	co.Conn = c
	if err = co.WriteMsg(m); err != nil {
		return nil, err
	}
	r, err = co.ReadMsg()
	return r, err
}

// Exchange performs an synchronous query. It sends the message m to the address
// contained in a and waits for an reply. Basic use pattern with a *dns.Client:
//
//	c := new(dns.Client)
//	in, rtt, err := c.Exchange(message, "127.0.0.1:53")
//
// Exchange does not retry a failed query, nor will it fall back to TCP in
// case of truncation.
func (c *Client) Exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	if !c.SingleInflight {
		return c.exchange(m, a)
	}
	// This adds a bunch of garbage, TODO(miek).
	t := "nop"
	if t1, ok := TypeToString[m.Question[0].Qtype]; ok {
		t = t1
	}
	cl := "nop"
	if cl1, ok := ClassToString[m.Question[0].Qclass]; ok {
		cl = cl1
	}
	r, rtt, err, shared := c.group.Do(m.Question[0].Name+t+cl, func() (*Msg, time.Duration, error) {
		return c.exchange(m, a)
	})
	if err != nil {
		return r, rtt, err
	}
	if shared {
		return r.Copy(), rtt, nil
	}
	return r, rtt, nil
}

func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	timeout := dnsTimeout
	var co *Conn
	if c.DialTimeout != 0 {
		timeout = c.DialTimeout
	}
	if c.Net == "" {
		co, err = DialTimeout("udp", a, timeout)
	} else {
		co, err = DialTimeout(c.Net, a, timeout)
	}
	if err != nil {
		return nil, 0, err
	}
	timeout = dnsTimeout
	if c.ReadTimeout != 0 {
		timeout = c.ReadTimeout
	}
	co.SetReadDeadline(time.Now().Add(timeout))
	timeout = dnsTimeout
	if c.WriteTimeout != 0 {
		timeout = c.WriteTimeout
	}
	co.SetWriteDeadline(time.Now().Add(timeout))
	defer co.Close()
	opt := m.IsEdns0()
	// If EDNS0 is used use that for size.
	if opt != nil && opt.UDPSize() >= MinMsgSize {
		co.UDPSize = opt.UDPSize()
	}
	// Otherwise use the client's configured UDP size.
	if opt == nil && c.UDPSize >= MinMsgSize {
		co.UDPSize = c.UDPSize
	}
	co.TsigSecret = c.TsigSecret
	if err = co.WriteMsg(m); err != nil {
		return nil, 0, err
	}
	r, err = co.ReadMsg()
	return r, co.rtt, err
}

// ReadMsg reads a message from the connection co.
// If the received message contains a TSIG record the transaction
// signature is verified.
func (co *Conn) ReadMsg() (*Msg, error) {
	var p []byte
	m := new(Msg)
	if _, ok := co.Conn.(*net.TCPConn); ok {
		p = make([]byte, MaxMsgSize)
	} else {
		if co.UDPSize >= 512 {
			p = make([]byte, co.UDPSize)
		} else {
			p = make([]byte, MinMsgSize)
		}
	}
	n, err := co.Read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	co.rtt = time.Since(co.t)
	if t := m.IsTsig(); t != nil {
		if _, ok := co.TsigSecret[t.Hdr.Name]; !ok {
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		err = TsigVerify(p, co.TsigSecret[t.Hdr.Name], co.tsigRequestMAC, false)
	}
	return m, err
}

// Read implements the net.Conn read method.
func (co *Conn) Read(p []byte) (n int, err error) {
	if co.Conn == nil {
		return 0, ErrConnEmpty
	}
	if len(p) < 2 {
		return 0, io.ErrShortBuffer
	}
	if t, ok := co.Conn.(*net.TCPConn); ok {
		n, err = t.Read(p[0:2])
		if err != nil || n != 2 {
			return n, err
		}
		l, _ := unpackUint16(p[0:2], 0)
		if l == 0 {
			return 0, ErrShortRead
		}
		if int(l) > len(p) {
			return int(l), io.ErrShortBuffer
		}
		n, err = t.Read(p[:l])
		if err != nil {
			return n, err
		}
		i := n
		for i < int(l) {
			j, err := t.Read(p[i:int(l)])
			if err != nil {
				return i, err
			}
			i += j
		}
		n = i
		return n, err
	}
	// UDP connection
	n, err = co.Conn.Read(p)
	if err != nil {
		return n, err
	}
	return n, err
}

// WriteMsg sends a message throught the connection co.
// If the message m contains a TSIG record the transaction
// signature is calculated.
func (co *Conn) WriteMsg(m *Msg) (err error) {
	var out []byte
	if t := m.IsTsig(); t != nil {
		mac := ""
		if _, ok := co.TsigSecret[t.Hdr.Name]; !ok {
			return ErrSecret
		}
		out, mac, err = TsigGenerate(m, co.TsigSecret[t.Hdr.Name], co.tsigRequestMAC, false)
		// Set for the next read, allthough only used in zone transfers
		co.tsigRequestMAC = mac
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	co.t = time.Now()
	if _, err = co.Write(out); err != nil {
		return err
	}
	return nil
}

// Write implements the net.Conn Write method.
func (co *Conn) Write(p []byte) (n int, err error) {
	if t, ok := co.Conn.(*net.TCPConn); ok {
		lp := len(p)
		if lp < 2 {
			return 0, io.ErrShortBuffer
		}
		if lp > MaxMsgSize {
			return 0, &Error{err: "message too large"}
		}
		l := make([]byte, 2, lp+2)
		l[0], l[1] = packUint16(uint16(lp))
		p = append(l, p...)
		n, err := io.Copy(t, bytes.NewReader(p))
		return int(n), err
	}
	n, err = co.Conn.(*net.UDPConn).Write(p)
	return n, err
}

// Dial connects to the address on the named network.
func Dial(network, address string) (conn *Conn, err error) {
	conn = new(Conn)
	conn.Conn, err = net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return conn, nil
}

// DialTimeout acts like Dial but takes a timeout.
func DialTimeout(network, address string, timeout time.Duration) (conn *Conn, err error) {
	conn = new(Conn)
	conn.Conn, err = net.DialTimeout(network, address, timeout)
	if err != nil {
		return nil, err
	}
	return conn, nil
}

// Close implements the net.Conn Close method.
func (co *Conn) Close() error { return co.Conn.Close() }

// LocalAddr implements the net.Conn LocalAddr method.
func (co *Conn) LocalAddr() net.Addr { return co.Conn.LocalAddr() }

// RemoteAddr implements the net.Conn RemoteAddr method.
func (co *Conn) RemoteAddr() net.Addr { return co.Conn.RemoteAddr() }

// SetDeadline implements the net.Conn SetDeadline method.
func (co *Conn) SetDeadline(t time.Time) error { return co.Conn.SetDeadline(t) }

// SetReadDeadline implements the net.Conn SetReadDeadline method.
func (co *Conn) SetReadDeadline(t time.Time) error { return co.Conn.SetReadDeadline(t) }

// SetWriteDeadline implements the net.Conn SetWriteDeadline method.
func (co *Conn) SetWriteDeadline(t time.Time) error { return co.Conn.SetWriteDeadline(t) }
<|endoftext|>


<commit_before>package tunnel

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/hashicorp/yamux"
	"github.com/koding/logging"
)

// Client is responsible for creating a control connection to a tunnel server,
// creating new tunnels and proxy them to tunnel server.
type Client struct {
	// underlying yamux session
	session *yamux.Session

	// config holds the ClientConfig
	config *ClientConfig

	// yamuxConfig is passed to new yamux.Session's
	yamuxConfig *yamux.Config
	log         logging.Logger

	mu          sync.Mutex // guards the following
	closed      bool       // if client calls Close() and quits
	startNotify chan bool  // notifies if client established a conn to server

	reqWg sync.WaitGroup

	// redialBackoff is used to reconnect in exponential backoff intervals
	redialBackoff backoff.BackOff
}

// ClientConfig defines the configuration for the Client
type ClientConfig struct {
	// Identifier is the secret token that needs to be passed to the server.
	// Required if FetchIdentifier is not set
	Identifier string

	// FetchIdentifier can be used to fetch identifier. Required if Identifier
	// is not set.
	FetchIdentifier func() (string, error)

	// ServerAddr defines the TCP address of the tunnel server to be connected. This is required.
	ServerAddr string

	// LocalAddr defines the TCP address of the local server. This is optional
	// if you want to specify a single TCP address. Otherwise the client will
	// always proxy to 127.0.0.1:incomingPort, where incomingPort is the
	// tunnelserver's public exposed Port.
	LocalAddr string

	// Debug enables debug mode, enable only if you want to debug the server.
	Debug bool

	// Log defines the logger. If nil a default logging.Logger is used.
	Log logging.Logger

	// YamuxConfig defines the config which passed to every new yamux.Session. If nil
	// yamux.DefaultConfig() is used.
	YamuxConfig *yamux.Config
}

// verify is used to verify the ClientConfig
func (c *ClientConfig) verify() error {
	if c.ServerAddr == "" {
		return errors.New("config.ServerAddr must be set")
	}

	if c.Identifier == "" && c.FetchIdentifier == nil {
		return errors.New("neither config.Identifier nor config.FetchIdentifier is set")
	}

	if c.YamuxConfig != nil {
		if err := yamux.VerifyConfig(c.YamuxConfig); err != nil {
			return err
		}
	}

	return nil
}

// NewClient creates a new tunnel that is established between the serverAddr
// and localAddr. It exits if it can't create a new control connection to the
// server. If localAddr is empty client will always try to proxy to a local
// port.
func NewClient(cfg *ClientConfig) (*Client, error) {
	yamuxConfig := yamux.DefaultConfig()
	if cfg.YamuxConfig != nil {
		yamuxConfig = cfg.YamuxConfig
	}

	log := newLogger("tunnel-client", cfg.Debug)
	if cfg.Log != nil {
		log = cfg.Log
	}

	if err := cfg.verify(); err != nil {
		return nil, err
	}

	forever := backoff.NewExponentialBackOff()
	forever.MaxElapsedTime = 365 * 24 * time.Hour // 1 year

	client := &Client{
		config:        cfg,
		log:           log,
		yamuxConfig:   yamuxConfig,
		redialBackoff: forever,
		startNotify:   make(chan bool, 1),
	}

	return client, nil
}

// Start starts the client and connects to the server with the identifier.
// client.FetchIdentifier() will be used if it's not nil. It's supports
// reconnecting with exponential backoff intervals when the connection to the
// server disconnects. Call client.Close() to shutdown the client completely. A
// successfull connection will cause StartNotify() to receive a value.
func (c *Client) Start() {
	id := func() (string, error) {
		if c.config.FetchIdentifier != nil {
			return c.config.FetchIdentifier()
		}

		return c.config.Identifier, nil
	}

	c.redialBackoff.Reset()
	for {
		time.Sleep(c.redialBackoff.NextBackOff())
		identifier, err := id()
		if err != nil {
			c.log.Critical("client fetch identifier err: %s", err.Error())
			continue
		}

		if err := c.connect(identifier); err != nil {
			c.log.Critical("client connect err: %s", err.Error())
		}

		// exit if closed
		c.mu.Lock()
		if c.closed {
			c.mu.Unlock()
			return
		}
		c.mu.Unlock()
	}
}

// StartNotify returns a channel that receives a single value when the client
// established a successfull connection to the server.
func (c *Client) StartNotify() <-chan bool {
	return c.startNotify
}

// Close closes the client and shutdowns the connection to the tunnel server
func (c *Client) Close() error {
	if c.session == nil {
		return errors.New("session is not initialized")
	}

	if err := c.session.GoAway(); err != nil {
		return err
	}

	c.mu.Lock()
	c.closed = true
	c.mu.Unlock()

	c.reqWg.Wait() // wait until all connections are finished
	if err := c.session.Close(); err != nil {
		return err
	}

	return nil
}

func (c *Client) connect(identifier string) error {
	c.log.Debug("Trying to connect to '%s' with identifier '%s'", c.config.ServerAddr, identifier)
	conn, err := net.Dial("tcp", c.config.ServerAddr)
	if err != nil {
		return err
	}

	remoteAddr := fmt.Sprintf("http://%s%s", conn.RemoteAddr(), controlPath)
	c.log.Debug("CONNECT to '%s'", remoteAddr)
	req, err := http.NewRequest("CONNECT", remoteAddr, nil)
	if err != nil {
		return fmt.Errorf("CONNECT %s", err)
	}

	req.Header.Set(xKTunnelIdentifier, identifier)
	c.log.Debug("Writing request to TCP: %+v", req)
	if err := req.Write(conn); err != nil {
		return err
	}

	c.log.Debug("Reading response from TCP")
	resp, err := http.ReadResponse(bufio.NewReader(conn), req)
	if err != nil {
		return fmt.Errorf("read response %s", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 && resp.Status != connected {
		out, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}

		return fmt.Errorf("proxy server: %s. err: %s", resp.Status, string(out))
	}

	c.session, err = yamux.Client(conn, c.yamuxConfig)
	if err != nil {
		return err
	}

	var stream net.Conn

	openStream := func() error {
		// this is blocking until client opens a session to us
		stream, err = c.session.Open()
		return err
	}

	// if we don't receive anything from the server, we'll timeout
	select {
	case err := <-async(openStream):
		if err != nil {
			return err
		}
	case <-time.After(time.Second * 10):
		if stream != nil {
			stream.Close()
		}
		return errors.New("timeout opening session")
	}

	if _, err := stream.Write([]byte(ctHandshakeRequest)); err != nil {
		return err
	}

	buf := make([]byte, len(ctHandshakeResponse))
	if _, err := stream.Read(buf); err != nil {
		return err
	}

	if string(buf) != ctHandshakeResponse {
		return fmt.Errorf("handshake aborted. got: %s", string(buf))
	}

	ct := newControl(stream)
	c.log.Debug("client has started successfully.")
	c.redialBackoff.Reset() // we successfully connected, so we can reset the backoff

	c.mu.Lock()
	if c.startNotify != nil && !c.closed {
		select {
		case c.startNotify <- true:
		default:
			// reaching here is a race condition, because it indicates
			// startNotify has already a value. We panic because it's library
			// level problem that needs immediate attention
			panic("startNotify chan is already full")
		}
	}
	c.mu.Unlock()

	return c.listenControl(ct)
}

func (c *Client) listenControl(ct *control) error {
	for {
		var msg controlMsg
		err := ct.dec.Decode(&msg)
		if err != nil {
			c.reqWg.Wait() // wait until all requests are finished

			c.session.GoAway()
			c.session.Close()
			return fmt.Errorf("decode err: '%s'", err)
		}

		c.log.Debug("Received control msg %+v", msg)

		switch msg.Action {
		case requestClientSession:
			c.reqWg.Add(1)
			c.log.Debug("Received request to open a session to server")
			go func() {
				if err := c.proxy(msg.LocalPort); err != nil {
					c.log.Error("Proxy err between remote and local: '%s'\n", err)
				}
			}()
		}
	}
}

func (c *Client) proxy(port string) error {
	c.log.Debug("Opening a new stream from server session")
	remote, err := c.session.Open()
	if err != nil {
		return err
	}
	defer remote.Close()

	if port == "0" {
		port = "80"
	}

	localAddr := "127.0.0.1:" + port
	if c.config.LocalAddr != "" {
		localAddr = c.config.LocalAddr
	}

	c.log.Debug("Dialing local server %s", localAddr)
	local, err := net.Dial("tcp", localAddr)
	if err != nil {
		return err
	}

	// if the remote stream closes, we need a way to finish reading from local
	// server. If we don't set a timeout, it'll wait forever.
	local.SetReadDeadline(time.Now().Add(time.Second * 5))

	c.log.Debug("Starting to proxy between remote and local server")
	c.join(local, remote)
	c.log.Debug("Proxing between remote and local server finished")
	return err
}

func (c *Client) join(local, remote net.Conn) {
	transfer := func(dst, src net.Conn, done chan struct{}) {
		_, err := io.Copy(dst, src)
		if err != nil {
			c.log.Error("copy error: %s", err.Error())
		}

		switch s := src.(type) {
		case *net.TCPConn:
			// only client -> local connections are pure tcp conns
			if err := s.CloseRead(); err != nil {
				c.log.Error("CloseRead error: %s", err.Error())
			}
		default:
			if err := s.Close(); err != nil {
				c.log.Error("Close error: %s", err.Error())
			}
		}

		done <- struct{}{}
	}

	wait := make(chan struct{}, 2)

	go transfer(local, remote, wait)
	go transfer(remote, local, wait)

	<-wait
	<-wait

	c.reqWg.Done()
}
<commit_msg>client.go: more improvements, still working on sync on connections<commit_after>package tunnel

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/hashicorp/yamux"
	"github.com/koding/logging"
)

// Client is responsible for creating a control connection to a tunnel server,
// creating new tunnels and proxy them to tunnel server.
type Client struct {
	// underlying yamux session
	session *yamux.Session

	// config holds the ClientConfig
	config *ClientConfig

	// yamuxConfig is passed to new yamux.Session's
	yamuxConfig *yamux.Config
	log         logging.Logger

	mu          sync.Mutex // guards the following
	closed      bool       // if client calls Close() and quits
	startNotify chan bool  // notifies if client established a conn to server

	reqWg sync.WaitGroup

	// redialBackoff is used to reconnect in exponential backoff intervals
	redialBackoff backoff.BackOff
}

// ClientConfig defines the configuration for the Client
type ClientConfig struct {
	// Identifier is the secret token that needs to be passed to the server.
	// Required if FetchIdentifier is not set
	Identifier string

	// FetchIdentifier can be used to fetch identifier. Required if Identifier
	// is not set.
	FetchIdentifier func() (string, error)

	// ServerAddr defines the TCP address of the tunnel server to be connected. This is required.
	ServerAddr string

	// LocalAddr defines the TCP address of the local server. This is optional
	// if you want to specify a single TCP address. Otherwise the client will
	// always proxy to 127.0.0.1:incomingPort, where incomingPort is the
	// tunnelserver's public exposed Port.
	LocalAddr string

	// Debug enables debug mode, enable only if you want to debug the server.
	Debug bool

	// Log defines the logger. If nil a default logging.Logger is used.
	Log logging.Logger

	// YamuxConfig defines the config which passed to every new yamux.Session. If nil
	// yamux.DefaultConfig() is used.
	YamuxConfig *yamux.Config
}

// verify is used to verify the ClientConfig
func (c *ClientConfig) verify() error {
	if c.ServerAddr == "" {
		return errors.New("config.ServerAddr must be set")
	}

	if c.Identifier == "" && c.FetchIdentifier == nil {
		return errors.New("neither config.Identifier nor config.FetchIdentifier is set")
	}

	if c.YamuxConfig != nil {
		if err := yamux.VerifyConfig(c.YamuxConfig); err != nil {
			return err
		}
	}

	return nil
}

// NewClient creates a new tunnel that is established between the serverAddr
// and localAddr. It exits if it can't create a new control connection to the
// server. If localAddr is empty client will always try to proxy to a local
// port.
func NewClient(cfg *ClientConfig) (*Client, error) {
	yamuxConfig := yamux.DefaultConfig()
	if cfg.YamuxConfig != nil {
		yamuxConfig = cfg.YamuxConfig
	}

	log := newLogger("tunnel-client", cfg.Debug)
	if cfg.Log != nil {
		log = cfg.Log
	}

	if err := cfg.verify(); err != nil {
		return nil, err
	}

	forever := backoff.NewExponentialBackOff()
	forever.MaxElapsedTime = 365 * 24 * time.Hour // 1 year

	client := &Client{
		config:        cfg,
		log:           log,
		yamuxConfig:   yamuxConfig,
		redialBackoff: forever,
		startNotify:   make(chan bool, 1),
	}

	return client, nil
}

// Start starts the client and connects to the server with the identifier.
// client.FetchIdentifier() will be used if it's not nil. It's supports
// reconnecting with exponential backoff intervals when the connection to the
// server disconnects. Call client.Close() to shutdown the client completely. A
// successfull connection will cause StartNotify() to receive a value.
func (c *Client) Start() {
	id := func() (string, error) {
		if c.config.FetchIdentifier != nil {
			return c.config.FetchIdentifier()
		}

		return c.config.Identifier, nil
	}

	c.redialBackoff.Reset()
	for {
		time.Sleep(c.redialBackoff.NextBackOff())
		identifier, err := id()
		if err != nil {
			c.log.Critical("client fetch identifier err: %s", err.Error())
			continue
		}

		if err := c.connect(identifier); err != nil {
			c.log.Critical("client connect err: %s", err.Error())
		}

		// exit if closed
		c.mu.Lock()
		if c.closed {
			c.mu.Unlock()
			return
		}
		c.mu.Unlock()
	}
}

// StartNotify returns a channel that receives a single value when the client
// established a successfull connection to the server.
func (c *Client) StartNotify() <-chan bool {
	return c.startNotify
}

// Close closes the client and shutdowns the connection to the tunnel server
func (c *Client) Close() error {
	if c.session == nil {
		return errors.New("session is not initialized")
	}

	if err := c.session.GoAway(); err != nil {
		return err
	}

	c.mu.Lock()
	c.closed = true
	c.mu.Unlock()

	c.reqWg.Wait() // wait until all connections are finished
	if err := c.session.Close(); err != nil {
		return err
	}

	return nil
}

func (c *Client) connect(identifier string) error {
	c.log.Debug("Trying to connect to '%s' with identifier '%s'", c.config.ServerAddr, identifier)
	conn, err := net.Dial("tcp", c.config.ServerAddr)
	if err != nil {
		return err
	}

	remoteAddr := fmt.Sprintf("http://%s%s", conn.RemoteAddr(), controlPath)
	c.log.Debug("CONNECT to '%s'", remoteAddr)
	req, err := http.NewRequest("CONNECT", remoteAddr, nil)
	if err != nil {
		return fmt.Errorf("CONNECT %s", err)
	}

	req.Header.Set(xKTunnelIdentifier, identifier)
	c.log.Debug("Writing request to TCP: %+v", req)
	if err := req.Write(conn); err != nil {
		return err
	}

	c.log.Debug("Reading response from TCP")
	resp, err := http.ReadResponse(bufio.NewReader(conn), req)
	if err != nil {
		return fmt.Errorf("read response %s", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 && resp.Status != connected {
		out, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}

		return fmt.Errorf("proxy server: %s. err: %s", resp.Status, string(out))
	}

	c.session, err = yamux.Client(conn, c.yamuxConfig)
	if err != nil {
		return err
	}

	var stream net.Conn

	openStream := func() error {
		// this is blocking until client opens a session to us
		stream, err = c.session.Open()
		return err
	}

	// if we don't receive anything from the server, we'll timeout
	select {
	case err := <-async(openStream):
		if err != nil {
			return err
		}
	case <-time.After(time.Second * 10):
		if stream != nil {
			stream.Close()
		}
		return errors.New("timeout opening session")
	}

	if _, err := stream.Write([]byte(ctHandshakeRequest)); err != nil {
		return err
	}

	buf := make([]byte, len(ctHandshakeResponse))
	if _, err := stream.Read(buf); err != nil {
		return err
	}

	if string(buf) != ctHandshakeResponse {
		return fmt.Errorf("handshake aborted. got: %s", string(buf))
	}

	ct := newControl(stream)
	c.log.Debug("client has started successfully.")
	c.redialBackoff.Reset() // we successfully connected, so we can reset the backoff

	c.mu.Lock()
	if c.startNotify != nil && !c.closed {
		select {
		case c.startNotify <- true:
		default:
			// reaching here is a race condition, because it indicates
			// startNotify has already a value. We panic because it's library
			// level problem that needs immediate attention
			panic("startNotify chan is already full")
		}
	}
	c.mu.Unlock()

	return c.listenControl(ct)
}

func (c *Client) listenControl(ct *control) error {
	for {
		var msg controlMsg
		err := ct.dec.Decode(&msg)
		if err != nil {
			c.reqWg.Wait() // wait until all requests are finished

			c.session.GoAway()
			c.session.Close()
			return fmt.Errorf("decode err: '%s'", err)
		}

		c.log.Debug("Received control msg %+v", msg)

		switch msg.Action {
		case requestClientSession:
			c.reqWg.Add(1)
			c.log.Debug("Received request to open a session to server")
			go func() {
				if err := c.proxy(msg.LocalPort); err != nil {
					c.log.Error("Proxy err between remote and local: '%s'\n", err)
				}
			}()
		}
	}
}

func (c *Client) proxy(port string) error {
	c.log.Debug("Opening a new stream from server session")
	remote, err := c.session.Open()
	if err != nil {
		return err
	}
	defer remote.Close()

	if port == "0" {
		port = "80"
	}

	localAddr := "127.0.0.1:" + port
	if c.config.LocalAddr != "" {
		localAddr = c.config.LocalAddr
	}

	c.log.Debug("Dialing local server %s", localAddr)
	local, err := net.Dial("tcp", localAddr)
	if err != nil {
		return err
	}

	c.log.Debug("Starting to proxy between remote and local server")
	join(local, remote)
	c.reqWg.Done()
	c.log.Debug("Proxing between remote and local server finished")
	return err
}

func join(local, remote net.Conn) {
	transfer := func(dst, src net.Conn, done chan struct{}) {
		_, err := io.Copy(dst, src)
		if err != nil {
			// log.Printf("copy error: %s\n", err.Error())
		}

		if err := src.Close(); err != nil {
			log.Printf("close error: %s\n", err.Error())
		}

		switch s := src.(type) {
		case *net.TCPConn:
			// only client -> local connections are pure tcp conns
			if err := s.CloseRead(); err != nil {
				log.Printf("closeRead error: %s\n", err.Error())
			}
		default:
			if err := s.Close(); err != nil {
				log.Printf("close error: %s\n", err.Error())
			}
		}

		done <- struct{}{}
	}

	remoteClosed := make(chan struct{}, 1)
	localClosed := make(chan struct{}, 1)

	go transfer(local, remote, remoteClosed)
	go transfer(remote, local, localClosed)

	// close other side of connections when we are done with it
	select {
	case <-localClosed:
		remote.Close()
	case <-remoteClosed:
		local.Close()
	}

	return
}
<|endoftext|>