<commit_before>package rtmp

import (
	"crypto/tls"
	"github.com/elobuff/goamf"
	"net"
	"net/url"
	"sync"
	"sync/atomic"
	"time"
)

type Client struct {
	url string

	connected bool

	conn net.Conn

	outBytes        uint32
	outMessages     chan *Message
	outWindowSize   uint32
	outChunkSize    uint32
	outChunkStreams map[uint32]*OutboundChunkStream
	inBytes         uint32
	inMessages      chan *Message
	inNotify        chan uint8
	inWindowSize    uint32
	inChunkSize     uint32
	inChunkStreams  map[uint32]*InboundChunkStream

	responses      map[uint32]*Response
	responsesMutex sync.Mutex

	lastTransactionId uint32
	connectionId      string

	amfExternalHandlers map[string]amf.ExternalHandler
}

func NewClient(url string) (c *Client) {
	c = &Client{
		url:                 url,
		amfExternalHandlers: make(map[string]amf.ExternalHandler),
	}

	c.Reset()

	return
}

func (c *Client) IsAlive() bool {
	if c.connected != true {
		return false
	}

	return true
}

func (c *Client) Reset() {
	c.connected = false

	if c.conn != nil {
		c.conn.Close()
	}

	if c.outMessages != nil {
		close(c.outMessages)
	}

	if c.inMessages != nil {
		close(c.inMessages)
	}

	c.outBytes = 0
	c.outMessages = make(chan *Message, 100)
	c.outChunkSize = DEFAULT_CHUNK_SIZE
	c.outWindowSize = DEFAULT_WINDOW_SIZE
	c.outChunkStreams = make(map[uint32]*OutboundChunkStream)
	c.inBytes = 0
	c.inMessages = make(chan *Message, 100)
	c.inChunkSize = DEFAULT_CHUNK_SIZE
	c.inWindowSize = DEFAULT_WINDOW_SIZE
	c.inChunkStreams = make(map[uint32]*InboundChunkStream)
	c.responses = make(map[uint32]*Response)
	c.lastTransactionId = 0
	c.connectionId = ""
}

func (c *Client) RegisterExternalHandler(name string, fn amf.ExternalHandler) {
	c.amfExternalHandlers[name] = fn
}

func (c *Client) Disconnect() {
	c.Reset()
	log.Info("disconnected from %s", c.url)
}

func (c *Client) Connect() (err error) {
	log.Info("connecting to %s", c.url)

	url, err := url.Parse(c.url)
	if err != nil {
		return err
	}

	switch url.Scheme {
	case "rtmp":
		c.conn, err = net.DialTimeout("tcp", url.Host, 5*time.Second)
		if err != nil {
			return err
		}
	case "rtmps":
		var nc net.Conn
		nc, err = net.DialTimeout("tcp", url.Host, 5*time.Second)
		if err != nil {
			return err
		}

		config := &tls.Config{InsecureSkipVerify: true}
		tc := tls.Client(nc, config)
		err = tc.Handshake()
		if err != nil {
			return err
		}

		c.conn = tc
	default:
		return Error("Unsupported scheme: %s", url.Scheme)
	}

	log.Debug("handshaking with %s", c.url)

	err = c.handshake()
	if err != nil {
		return Error("client connect: could not complete handshake: %s", err)
	}

	log.Debug("sending connect command to %s", c.url)

	go c.receiveLoop()
	go c.sendLoop()
	go c.routeLoop()

	var id string
	id, err = c.connect()
	if err != nil {
		return Error("client connect: could not complete connect: %s", err)
	}

	c.connected = true
	c.connectionId = id

	log.Info("connected to %s (%s)", c.url, c.connectionId)

	return
}

func (c *Client) NextTransactionId() uint32 {
	return atomic.AddUint32(&c.lastTransactionId, 1)
}

func (c *Client) GetResponse(tid uint32) (response *Response, ready bool) {
	c.responsesMutex.Lock()
	defer c.responsesMutex.Unlock()
	response = c.responses[tid]
	if response != nil {
		ready = true
		delete(c.responses, tid)
	}
	return
}

func (c *Client) SendMessage(msg *Message) {
	log.Debug("SENDING MESSAGE: %d", msg.TransactionId)
	c.outMessages <- msg
}

func (c *Client) Call(msg *Message, t uint32) (response *Response, err error) {
	c.SendMessage(msg)

	tid := msg.TransactionId
	poll := time.Tick(time.Duration(5) * time.Millisecond)
	timeout := time.After(time.Duration(t) * time.Second)

	for {
		select {
		case <-poll:
			c.responsesMutex.Lock()
			response = c.responses[tid]
			if response != nil {
				log.Trace("client call: found response for %d", tid)
				delete(c.responses, tid)
			}
			c.responsesMutex.Unlock()

			if response != nil {
				return
			}
		case <-timeout:
			return response, Error("timed out (no response after %d seconds)", t)
		}
	}

	return
}

func (c *Client) Read(p []byte) (n int, err error) {
	n, err = c.conn.Read(p)
	c.inBytes += uint32(n)
	log.Trace("read %d", n)
	return n, err
}

func (c *Client) Write(p []byte) (n int, err error) {
	n, err = c.conn.Write(p)
	c.outBytes += uint32(n)
	log.Trace("write %d", n)
	return n, err
}
<commit_msg>Decrease info level verbosity<commit_after>package rtmp

import (
	"crypto/tls"
	"github.com/elobuff/goamf"
	"net"
	"net/url"
	"sync"
	"sync/atomic"
	"time"
)

type Client struct {
	url string

	connected bool

	conn net.Conn

	outBytes        uint32
	outMessages     chan *Message
	outWindowSize   uint32
	outChunkSize    uint32
	outChunkStreams map[uint32]*OutboundChunkStream
	inBytes         uint32
	inMessages      chan *Message
	inNotify        chan uint8
	inWindowSize    uint32
	inChunkSize     uint32
	inChunkStreams  map[uint32]*InboundChunkStream

	responses      map[uint32]*Response
	responsesMutex sync.Mutex

	lastTransactionId uint32
	connectionId      string

	amfExternalHandlers map[string]amf.ExternalHandler
}

func NewClient(url string) (c *Client) {
	c = &Client{
		url:                 url,
		amfExternalHandlers: make(map[string]amf.ExternalHandler),
	}

	c.Reset()

	return
}

func (c *Client) IsAlive() bool {
	if c.connected != true {
		return false
	}

	return true
}

func (c *Client) Reset() {
	c.connected = false

	if c.conn != nil {
		c.conn.Close()
	}

	if c.outMessages != nil {
		close(c.outMessages)
	}

	if c.inMessages != nil {
		close(c.inMessages)
	}

	c.outBytes = 0
	c.outMessages = make(chan *Message, 100)
	c.outChunkSize = DEFAULT_CHUNK_SIZE
	c.outWindowSize = DEFAULT_WINDOW_SIZE
	c.outChunkStreams = make(map[uint32]*OutboundChunkStream)
	c.inBytes = 0
	c.inMessages = make(chan *Message, 100)
	c.inChunkSize = DEFAULT_CHUNK_SIZE
	c.inWindowSize = DEFAULT_WINDOW_SIZE
	c.inChunkStreams = make(map[uint32]*InboundChunkStream)
	c.responses = make(map[uint32]*Response)
	c.lastTransactionId = 0
	c.connectionId = ""
}

func (c *Client) RegisterExternalHandler(name string, fn amf.ExternalHandler) {
	c.amfExternalHandlers[name] = fn
}

func (c *Client) Disconnect() {
	c.Reset()
	log.Debug("disconnected from %s", c.url)
}

func (c *Client) Connect() (err error) {
	log.Debug("connecting to %s", c.url)

	url, err := url.Parse(c.url)
	if err != nil {
		return err
	}

	switch url.Scheme {
	case "rtmp":
		c.conn, err = net.DialTimeout("tcp", url.Host, 5*time.Second)
		if err != nil {
			return err
		}
	case "rtmps":
		var nc net.Conn
		nc, err = net.DialTimeout("tcp", url.Host, 5*time.Second)
		if err != nil {
			return err
		}

		config := &tls.Config{InsecureSkipVerify: true}
		tc := tls.Client(nc, config)
		err = tc.Handshake()
		if err != nil {
			return err
		}

		c.conn = tc
	default:
		return Error("Unsupported scheme: %s", url.Scheme)
	}

	log.Debug("handshaking with %s", c.url)

	err = c.handshake()
	if err != nil {
		return Error("client connect: could not complete handshake: %s", err)
	}

	log.Debug("sending connect command to %s", c.url)

	go c.receiveLoop()
	go c.sendLoop()
	go c.routeLoop()

	var id string
	id, err = c.connect()
	if err != nil {
		return Error("client connect: could not complete connect: %s", err)
	}

	c.connected = true
	c.connectionId = id

	log.Debug("connected to %s (%s)", c.url, c.connectionId)

	return
}

func (c *Client) NextTransactionId() uint32 {
	return atomic.AddUint32(&c.lastTransactionId, 1)
}

func (c *Client) GetResponse(tid uint32) (response *Response, ready bool) {
	c.responsesMutex.Lock()
	defer c.responsesMutex.Unlock()
	response = c.responses[tid]
	if response != nil {
		ready = true
		delete(c.responses, tid)
	}
	return
}

func (c *Client) SendMessage(msg *Message) {
	log.Debug("SENDING MESSAGE: %d", msg.TransactionId)
	c.outMessages <- msg
}

func (c *Client) Call(msg *Message, t uint32) (response *Response, err error) {
	c.SendMessage(msg)

	tid := msg.TransactionId
	poll := time.Tick(time.Duration(5) * time.Millisecond)
	timeout := time.After(time.Duration(t) * time.Second)

	for {
		select {
		case <-poll:
			c.responsesMutex.Lock()
			response = c.responses[tid]
			if response != nil {
				log.Trace("client call: found response for %d", tid)
				delete(c.responses, tid)
			}
			c.responsesMutex.Unlock()

			if response != nil {
				return
			}
		case <-timeout:
			return response, Error("timed out (no response after %d seconds)", t)
		}
	}

	return
}

func (c *Client) Read(p []byte) (n int, err error) {
	n, err = c.conn.Read(p)
	c.inBytes += uint32(n)
	log.Trace("read %d", n)
	return n, err
}

func (c *Client) Write(p []byte) (n int, err error) {
	n, err = c.conn.Write(p)
	c.outBytes += uint32(n)
	log.Trace("write %d", n)
	return n, err
}
<|endoftext|>


<commit_before>package atlas // import "github.com/keltia/ripe-atlas"

import (
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/keltia/proxy"
	"github.com/pkg/errors"
)

// NewClient is the first function to call.
// Yes, it does take multiple config
// and the last one wins.
func NewClient(cfgs ...Config) (*Client, error) {
	c := &Client{}
	for _, cfg := range cfgs {
		c.config = cfg
	}

	// This holds the global options
	c.opts = make(map[string]string)

	// If no log output is specified, use the default one
	if c.config.Log == nil {
		c.log = log.New(os.Stderr, "", log.LstdFlags|log.LUTC)
	} else {
		c.log = c.config.Log
	}

	// Ensure this is not empty
	if c.config.endpoint == "" {
		c.config.endpoint = apiEndpoint
	}
	c.verbose("c.config=%#v", c.config)

	// Create and save the http.Client
	return c.addHTTPClient()
}

// HasAPIKey returns whether an API key is stored
func (c *Client) HasAPIKey() (string, bool) {
	if c.config.APIKey == "" {
		return "", false
	}
	return c.config.APIKey, true
}

// call is s shortcut
func (c *Client) call(req *http.Request) (*http.Response, error) {
	c.verbose("Full URL:\n%v", req.URL)

	myurl, _ := url.Parse(apiEndpoint)
	req.Header.Set("Host", myurl.Host)
	req.Header.Set("User-Agent", fmt.Sprintf("ripe-atlas/%s", ourVersion))

	return c.client.Do(req)
}

func (c *Client) addHTTPClient() (*Client, error) {
	_, transport := proxy.SetupTransport(apiEndpoint)
	if transport == nil {
		return c, errors.New("addhttpclient")
	}
	c.client = &http.Client{Transport: transport, Timeout: 20 * time.Second}
	return c, nil
}

// SetOption sets a global option
func (c *Client) SetOption(name, value string) *Client {
	if value != "" {
		c.opts[name] = value
	}
	return c
}

// GetVersion returns the API wrapper version
func GetVersion() string {
	return ourVersion
}
<commit_msg>Handle the logging levels.<commit_after>package atlas // import "github.com/keltia/ripe-atlas"

import (
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/keltia/proxy"
	"github.com/pkg/errors"
)

// NewClient is the first function to call.
// Yes, it does take multiple config
// and the last one wins.
func NewClient(cfgs ...Config) (*Client, error) {
	c := &Client{}
	for _, cfg := range cfgs {
		c.config = cfg
	}

	// This holds the global options
	c.opts = make(map[string]string)

	// If no log output is specified, use the default one
	if c.config.Log == nil {
		c.log = log.New(os.Stderr, "", log.LstdFlags|log.LUTC)
	} else {
		c.log = c.config.Log
	}

	// Set log levels
	if c.config.Verbose == true {
		c.level = 1
	}

	if c.config.Level > 2 {
		c.level = 2
	}

	// Ensure this is not empty
	if c.config.endpoint == "" {
		c.config.endpoint = apiEndpoint
	}
	c.verbose("c.config=%#v", c.config)

	// Create and save the http.Client
	return c.addHTTPClient()
}

// HasAPIKey returns whether an API key is stored
func (c *Client) HasAPIKey() (string, bool) {
	if c.config.APIKey == "" {
		return "", false
	}
	return c.config.APIKey, true
}

// call is s shortcut
func (c *Client) call(req *http.Request) (*http.Response, error) {
	c.verbose("Full URL:\n%v", req.URL)

	myurl, _ := url.Parse(apiEndpoint)
	req.Header.Set("Host", myurl.Host)
	req.Header.Set("User-Agent", fmt.Sprintf("ripe-atlas/%s", ourVersion))

	return c.client.Do(req)
}

func (c *Client) addHTTPClient() (*Client, error) {
	_, transport := proxy.SetupTransport(apiEndpoint)
	if transport == nil {
		return c, errors.New("addhttpclient")
	}
	c.client = &http.Client{Transport: transport, Timeout: 20 * time.Second}
	return c, nil
}

// SetOption sets a global option
func (c *Client) SetOption(name, value string) *Client {
	if value != "" {
		c.opts[name] = value
	}
	return c
}

// GetVersion returns the API wrapper version
func GetVersion() string {
	return ourVersion
}
<|endoftext|>


<commit_before>package libovsdb

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/url"
	"reflect"
	"strings"
	"sync"

	"github.com/cenkalti/rpc2"
	"github.com/cenkalti/rpc2/jsonrpc"
)

// OvsdbClient is an OVSDB client
type OvsdbClient struct {
	rpcClient     *rpc2.Client
	Schema        map[string]DatabaseSchema
	handlers      []NotificationHandler
	handlersMutex *sync.Mutex
}

func newOvsdbClient(c *rpc2.Client) *OvsdbClient {
	ovs := &OvsdbClient{
		rpcClient:     c,
		Schema:        make(map[string]DatabaseSchema),
		handlersMutex: &sync.Mutex{},
	}
	connectionsMutex.Lock()
	defer connectionsMutex.Unlock()
	if connections == nil {
		connections = make(map[*rpc2.Client]*OvsdbClient)
	}
	connections[c] = ovs
	return ovs
}

// Would rather replace this connection map with an OvsdbClient Receiver scoped method
// Unfortunately rpc2 package acts wierd with a receiver scoped method and needs some investigation.
var (
	connections      map[*rpc2.Client]*OvsdbClient
	connectionsMutex = &sync.RWMutex{}
)

// Constants defined for libovsdb
const (
	defaultTCPAddress  = "127.0.0.1:6640"
	defaultUnixAddress = "/var/run/openvswitch/ovnnb_db.sock"
	SSL                = "ssl"
	TCP                = "tcp"
	UNIX               = "unix"
)

// Connect to ovn, using endpoint in format ovsdb Connection Methods
// If address is empty, use default address for specified protocol
func Connect(endpoints string, tlsConfig *tls.Config) (*OvsdbClient, error) {
	var c net.Conn
	var err error
	var u *url.URL

	for _, endpoint := range strings.Split(endpoints, ",") {
		if u, err = url.Parse(endpoint); err != nil {
			return nil, err
		}
		// u.Opaque contains the original endPoint with the leading protocol stripped
		// off. For example: endPoint is "tcp:127.0.0.1:6640" and u.Opaque is "127.0.0.1:6640"
		host := u.Opaque
		if len(host) == 0 {
			host = defaultTCPAddress
		}
		switch u.Scheme {
		case UNIX:
			path := u.Path
			if len(path) == 0 {
				path = defaultUnixAddress
			}
			c, err = net.Dial(u.Scheme, path)
		case TCP:
			c, err = net.Dial(u.Scheme, host)
		case SSL:
			c, err = tls.Dial("tcp", host, tlsConfig)
		default:
			err = fmt.Errorf("unknown network protocol %s", u.Scheme)
		}

		if err == nil {
			return newRPC2Client(c)
		}
	}

	return nil, fmt.Errorf("failed to connect: %s", err.Error())
}

func newRPC2Client(conn net.Conn) (*OvsdbClient, error) {
	c := rpc2.NewClientWithCodec(jsonrpc.NewJSONCodec(conn))
	c.SetBlocking(true)
	c.Handle("echo", echo)
	c.Handle("update", update)
	go c.Run()
	go handleDisconnectNotification(c)

	ovs := newOvsdbClient(c)

	// Process Async Notifications
	dbs, err := ovs.ListDbs()
	if err != nil {
		return nil, err
	}

	for _, db := range dbs {
		schema, err := ovs.GetSchema(db)
		if err == nil {
			ovs.Schema[db] = *schema
		} else {
			return nil, err
		}
	}
	return ovs, nil
}

// Register registers the supplied NotificationHandler to recieve OVSDB Notifications
func (ovs *OvsdbClient) Register(handler NotificationHandler) {
	ovs.handlersMutex.Lock()
	defer ovs.handlersMutex.Unlock()
	ovs.handlers = append(ovs.handlers, handler)
}

//Get Handler by index
func getHandlerIndex(handler NotificationHandler, handlers []NotificationHandler) (int, error) {
	for i, h := range handlers {
		if reflect.DeepEqual(h, handler) {
			return i, nil
		}
	}
	return -1, errors.New("Handler not found")
}

// Unregister the supplied NotificationHandler to not recieve OVSDB Notifications anymore
func (ovs *OvsdbClient) Unregister(handler NotificationHandler) error {
	ovs.handlersMutex.Lock()
	defer ovs.handlersMutex.Unlock()
	i, err := getHandlerIndex(handler, ovs.handlers)
	if err != nil {
		return err
	}
	ovs.handlers = append(ovs.handlers[:i], ovs.handlers[i+1:]...)
	return nil
}

// NotificationHandler is the interface that must be implemented to receive notifcations
type NotificationHandler interface {
	// RFC 7047 section 4.1.6 Update Notification
	Update(context interface{}, tableUpdates TableUpdates)

	// RFC 7047 section 4.1.9 Locked Notification
	Locked([]interface{})

	// RFC 7047 section 4.1.10 Stolen Notification
	Stolen([]interface{})

	// RFC 7047 section 4.1.11 Echo Notification
	Echo([]interface{})

	Disconnected(*OvsdbClient)
}

// RFC 7047 : Section 4.1.6 : Echo
func echo(client *rpc2.Client, args []interface{}, reply *[]interface{}) error {
	*reply = args
	connectionsMutex.RLock()
	defer connectionsMutex.RUnlock()
	if _, ok := connections[client]; ok {
		connections[client].handlersMutex.Lock()
		defer connections[client].handlersMutex.Unlock()
		for _, handler := range connections[client].handlers {
			handler.Echo(nil)
		}
	}
	return nil
}

// RFC 7047 : Update Notification Section 4.1.6
// Processing "params": [<json-value>, <table-updates>]
func update(client *rpc2.Client, params []interface{}, reply *interface{}) error {
	if len(params) < 2 {
		return errors.New("Invalid Update message")
	}
	// Ignore params[0] as we dont use the <json-value> currently for comparison

	raw, ok := params[1].(map[string]interface{})
	if !ok {
		return errors.New("Invalid Update message")
	}
	var rowUpdates map[string]map[string]RowUpdate

	b, err := json.Marshal(raw)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &rowUpdates)
	if err != nil {
		return err
	}

	// Update the local DB cache with the tableUpdates
	tableUpdates := getTableUpdatesFromRawUnmarshal(rowUpdates)
	connectionsMutex.RLock()
	defer connectionsMutex.RUnlock()
	if _, ok := connections[client]; ok {
		connections[client].handlersMutex.Lock()
		defer connections[client].handlersMutex.Unlock()
		for _, handler := range connections[client].handlers {
			handler.Update(params[0], tableUpdates)
		}
	}

	return nil
}

// GetSchema returns the schema in use for the provided database name
// RFC 7047 : get_schema
func (ovs OvsdbClient) GetSchema(dbName string) (*DatabaseSchema, error) {
	args := NewGetSchemaArgs(dbName)
	var reply DatabaseSchema
	err := ovs.rpcClient.Call("get_schema", args, &reply)
	if err != nil {
		return nil, err
	}
	ovs.Schema[dbName] = reply
	return &reply, err
}

// ListDbs returns the list of databases on the server
// RFC 7047 : list_dbs
func (ovs OvsdbClient) ListDbs() ([]string, error) {
	var dbs []string
	err := ovs.rpcClient.Call("list_dbs", nil, &dbs)
	if err != nil {
		return nil, fmt.Errorf("ListDbs failure - %v", err)
	}
	return dbs, err
}

// Transact performs the provided Operation's on the database
// RFC 7047 : transact
func (ovs OvsdbClient) Transact(database string, operation ...Operation) ([]OperationResult, error) {
	var reply []OperationResult
	db, ok := ovs.Schema[database]
	if !ok {
		return nil, fmt.Errorf("invalid Database %q Schema", database)
	}

	if ok := db.validateOperations(operation...); !ok {
		return nil, errors.New("Validation failed for the operation")
	}

	args := NewTransactArgs(database, operation...)
	err := ovs.rpcClient.Call("transact", args, &reply)
	if err != nil {
		return nil, err
	}
	return reply, nil
}

// MonitorAll is a convenience method to monitor every table/column
func (ovs OvsdbClient) MonitorAll(database string, jsonContext interface{}) (*TableUpdates, error) {
	schema, ok := ovs.Schema[database]
	if !ok {
		return nil, fmt.Errorf("invalid Database %q Schema", database)
	}

	requests := make(map[string]MonitorRequest)
	for table, tableSchema := range schema.Tables {
		var columns []string
		for column := range tableSchema.Columns {
			columns = append(columns, column)
		}
		requests[table] = MonitorRequest{
			Columns: columns,
			Select: MonitorSelect{
				Initial: true,
				Insert:  true,
				Delete:  true,
				Modify:  true,
			}}
	}
	return ovs.Monitor(database, jsonContext, requests)
}

// MonitorCancel will request cancel a previously issued monitor request
// RFC 7047 : monitor_cancel
func (ovs OvsdbClient) MonitorCancel(database string, jsonContext interface{}) error {
	var reply OperationResult

	args := NewMonitorCancelArgs(jsonContext)

	err := ovs.rpcClient.Call("monitor_cancel", args, &reply)
	if err != nil {
		return err
	}
	if reply.Error != "" {
		return fmt.Errorf("Error while executing transaction: %s", reply.Error)
	}
	return nil
}

// Monitor will provide updates for a given table/column
// RFC 7047 : monitor
func (ovs OvsdbClient) Monitor(database string, jsonContext interface{}, requests map[string]MonitorRequest) (*TableUpdates, error) {
	var reply TableUpdates

	args := NewMonitorArgs(database, jsonContext, requests)

	// This totally sucks. Refer to golang JSON issue #6213
	var response map[string]map[string]RowUpdate
	err := ovs.rpcClient.Call("monitor", args, &response)
	reply = getTableUpdatesFromRawUnmarshal(response)
	if err != nil {
		return nil, err
	}
	return &reply, err
}

func getTableUpdatesFromRawUnmarshal(raw map[string]map[string]RowUpdate) TableUpdates {
	var tableUpdates TableUpdates
	tableUpdates.Updates = make(map[string]TableUpdate)
	for table, update := range raw {
		tableUpdate := TableUpdate{update}
		tableUpdates.Updates[table] = tableUpdate
	}
	return tableUpdates
}

func clearConnection(c *rpc2.Client) {
	connectionsMutex.Lock()
	defer connectionsMutex.Unlock()
	if _, ok := connections[c]; ok {
		for _, handler := range connections[c].handlers {
			if handler != nil {
				handler.Disconnected(connections[c])
			}
		}
	}
	delete(connections, c)
}

func handleDisconnectNotification(c *rpc2.Client) {
	disconnected := c.DisconnectNotify()
	select {
	case <-disconnected:
		clearConnection(c)
	}
}

// Disconnect will close the OVSDB connection
func (ovs OvsdbClient) Disconnect() {
	ovs.rpcClient.Close()
}
<commit_msg>cleanup the underlying rpc client on erros in newRPC2Client()<commit_after>package libovsdb

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/url"
	"reflect"
	"strings"
	"sync"

	"github.com/cenkalti/rpc2"
	"github.com/cenkalti/rpc2/jsonrpc"
)

// OvsdbClient is an OVSDB client
type OvsdbClient struct {
	rpcClient     *rpc2.Client
	Schema        map[string]DatabaseSchema
	handlers      []NotificationHandler
	handlersMutex *sync.Mutex
}

func newOvsdbClient(c *rpc2.Client) *OvsdbClient {
	ovs := &OvsdbClient{
		rpcClient:     c,
		Schema:        make(map[string]DatabaseSchema),
		handlersMutex: &sync.Mutex{},
	}
	return ovs
}

// Would rather replace this connection map with an OvsdbClient Receiver scoped method
// Unfortunately rpc2 package acts wierd with a receiver scoped method and needs some investigation.
var (
	connections      map[*rpc2.Client]*OvsdbClient
	connectionsMutex = &sync.RWMutex{}
)

// Constants defined for libovsdb
const (
	defaultTCPAddress  = "127.0.0.1:6640"
	defaultUnixAddress = "/var/run/openvswitch/ovnnb_db.sock"
	SSL                = "ssl"
	TCP                = "tcp"
	UNIX               = "unix"
)

// Connect to ovn, using endpoint in format ovsdb Connection Methods
// If address is empty, use default address for specified protocol
func Connect(endpoints string, tlsConfig *tls.Config) (*OvsdbClient, error) {
	var c net.Conn
	var err error
	var u *url.URL

	for _, endpoint := range strings.Split(endpoints, ",") {
		if u, err = url.Parse(endpoint); err != nil {
			return nil, err
		}
		// u.Opaque contains the original endPoint with the leading protocol stripped
		// off. For example: endPoint is "tcp:127.0.0.1:6640" and u.Opaque is "127.0.0.1:6640"
		host := u.Opaque
		if len(host) == 0 {
			host = defaultTCPAddress
		}
		switch u.Scheme {
		case UNIX:
			path := u.Path
			if len(path) == 0 {
				path = defaultUnixAddress
			}
			c, err = net.Dial(u.Scheme, path)
		case TCP:
			c, err = net.Dial(u.Scheme, host)
		case SSL:
			c, err = tls.Dial("tcp", host, tlsConfig)
		default:
			err = fmt.Errorf("unknown network protocol %s", u.Scheme)
		}

		if err == nil {
			return newRPC2Client(c)
		}
	}

	return nil, fmt.Errorf("failed to connect: %s", err.Error())
}

func newRPC2Client(conn net.Conn) (*OvsdbClient, error) {
	c := rpc2.NewClientWithCodec(jsonrpc.NewJSONCodec(conn))
	c.SetBlocking(true)
	c.Handle("echo", echo)
	c.Handle("update", update)
	go c.Run()
	go handleDisconnectNotification(c)

	ovs := newOvsdbClient(c)

	// Process Async Notifications
	dbs, err := ovs.ListDbs()
	if err != nil {
		c.Close()
		return nil, err
	}

	for _, db := range dbs {
		schema, err := ovs.GetSchema(db)
		if err == nil {
			ovs.Schema[db] = *schema
		} else {
			c.Close()
			return nil, err
		}
	}

	connectionsMutex.Lock()
	defer connectionsMutex.Unlock()
	if connections == nil {
		connections = make(map[*rpc2.Client]*OvsdbClient)
	}
	connections[c] = ovs
	return ovs, nil
}

// Register registers the supplied NotificationHandler to recieve OVSDB Notifications
func (ovs *OvsdbClient) Register(handler NotificationHandler) {
	ovs.handlersMutex.Lock()
	defer ovs.handlersMutex.Unlock()
	ovs.handlers = append(ovs.handlers, handler)
}

//Get Handler by index
func getHandlerIndex(handler NotificationHandler, handlers []NotificationHandler) (int, error) {
	for i, h := range handlers {
		if reflect.DeepEqual(h, handler) {
			return i, nil
		}
	}
	return -1, errors.New("Handler not found")
}

// Unregister the supplied NotificationHandler to not recieve OVSDB Notifications anymore
func (ovs *OvsdbClient) Unregister(handler NotificationHandler) error {
	ovs.handlersMutex.Lock()
	defer ovs.handlersMutex.Unlock()
	i, err := getHandlerIndex(handler, ovs.handlers)
	if err != nil {
		return err
	}
	ovs.handlers = append(ovs.handlers[:i], ovs.handlers[i+1:]...)
	return nil
}

// NotificationHandler is the interface that must be implemented to receive notifcations
type NotificationHandler interface {
	// RFC 7047 section 4.1.6 Update Notification
	Update(context interface{}, tableUpdates TableUpdates)

	// RFC 7047 section 4.1.9 Locked Notification
	Locked([]interface{})

	// RFC 7047 section 4.1.10 Stolen Notification
	Stolen([]interface{})

	// RFC 7047 section 4.1.11 Echo Notification
	Echo([]interface{})

	Disconnected(*OvsdbClient)
}

// RFC 7047 : Section 4.1.6 : Echo
func echo(client *rpc2.Client, args []interface{}, reply *[]interface{}) error {
	*reply = args
	connectionsMutex.RLock()
	defer connectionsMutex.RUnlock()
	if _, ok := connections[client]; ok {
		connections[client].handlersMutex.Lock()
		defer connections[client].handlersMutex.Unlock()
		for _, handler := range connections[client].handlers {
			handler.Echo(nil)
		}
	}
	return nil
}

// RFC 7047 : Update Notification Section 4.1.6
// Processing "params": [<json-value>, <table-updates>]
func update(client *rpc2.Client, params []interface{}, reply *interface{}) error {
	if len(params) < 2 {
		return errors.New("Invalid Update message")
	}
	// Ignore params[0] as we dont use the <json-value> currently for comparison

	raw, ok := params[1].(map[string]interface{})
	if !ok {
		return errors.New("Invalid Update message")
	}
	var rowUpdates map[string]map[string]RowUpdate

	b, err := json.Marshal(raw)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b, &rowUpdates)
	if err != nil {
		return err
	}

	// Update the local DB cache with the tableUpdates
	tableUpdates := getTableUpdatesFromRawUnmarshal(rowUpdates)
	connectionsMutex.RLock()
	defer connectionsMutex.RUnlock()
	if _, ok := connections[client]; ok {
		connections[client].handlersMutex.Lock()
		defer connections[client].handlersMutex.Unlock()
		for _, handler := range connections[client].handlers {
			handler.Update(params[0], tableUpdates)
		}
	}

	return nil
}

// GetSchema returns the schema in use for the provided database name
// RFC 7047 : get_schema
func (ovs OvsdbClient) GetSchema(dbName string) (*DatabaseSchema, error) {
	args := NewGetSchemaArgs(dbName)
	var reply DatabaseSchema
	err := ovs.rpcClient.Call("get_schema", args, &reply)
	if err != nil {
		return nil, err
	}
	ovs.Schema[dbName] = reply
	return &reply, err
}

// ListDbs returns the list of databases on the server
// RFC 7047 : list_dbs
func (ovs OvsdbClient) ListDbs() ([]string, error) {
	var dbs []string
	err := ovs.rpcClient.Call("list_dbs", nil, &dbs)
	if err != nil {
		return nil, fmt.Errorf("ListDbs failure - %v", err)
	}
	return dbs, err
}

// Transact performs the provided Operation's on the database
// RFC 7047 : transact
func (ovs OvsdbClient) Transact(database string, operation ...Operation) ([]OperationResult, error) {
	var reply []OperationResult
	db, ok := ovs.Schema[database]
	if !ok {
		return nil, fmt.Errorf("invalid Database %q Schema", database)
	}

	if ok := db.validateOperations(operation...); !ok {
		return nil, errors.New("Validation failed for the operation")
	}

	args := NewTransactArgs(database, operation...)
	err := ovs.rpcClient.Call("transact", args, &reply)
	if err != nil {
		return nil, err
	}
	return reply, nil
}

// MonitorAll is a convenience method to monitor every table/column
func (ovs OvsdbClient) MonitorAll(database string, jsonContext interface{}) (*TableUpdates, error) {
	schema, ok := ovs.Schema[database]
	if !ok {
		return nil, fmt.Errorf("invalid Database %q Schema", database)
	}

	requests := make(map[string]MonitorRequest)
	for table, tableSchema := range schema.Tables {
		var columns []string
		for column := range tableSchema.Columns {
			columns = append(columns, column)
		}
		requests[table] = MonitorRequest{
			Columns: columns,
			Select: MonitorSelect{
				Initial: true,
				Insert:  true,
				Delete:  true,
				Modify:  true,
			}}
	}
	return ovs.Monitor(database, jsonContext, requests)
}

// MonitorCancel will request cancel a previously issued monitor request
// RFC 7047 : monitor_cancel
func (ovs OvsdbClient) MonitorCancel(database string, jsonContext interface{}) error {
	var reply OperationResult

	args := NewMonitorCancelArgs(jsonContext)

	err := ovs.rpcClient.Call("monitor_cancel", args, &reply)
	if err != nil {
		return err
	}
	if reply.Error != "" {
		return fmt.Errorf("Error while executing transaction: %s", reply.Error)
	}
	return nil
}

// Monitor will provide updates for a given table/column
// RFC 7047 : monitor
func (ovs OvsdbClient) Monitor(database string, jsonContext interface{}, requests map[string]MonitorRequest) (*TableUpdates, error) {
	var reply TableUpdates

	args := NewMonitorArgs(database, jsonContext, requests)

	// This totally sucks. Refer to golang JSON issue #6213
	var response map[string]map[string]RowUpdate
	err := ovs.rpcClient.Call("monitor", args, &response)
	reply = getTableUpdatesFromRawUnmarshal(response)
	if err != nil {
		return nil, err
	}
	return &reply, err
}

func getTableUpdatesFromRawUnmarshal(raw map[string]map[string]RowUpdate) TableUpdates {
	var tableUpdates TableUpdates
	tableUpdates.Updates = make(map[string]TableUpdate)
	for table, update := range raw {
		tableUpdate := TableUpdate{update}
		tableUpdates.Updates[table] = tableUpdate
	}
	return tableUpdates
}

func clearConnection(c *rpc2.Client) {
	connectionsMutex.Lock()
	defer connectionsMutex.Unlock()
	if _, ok := connections[c]; ok {
		for _, handler := range connections[c].handlers {
			if handler != nil {
				handler.Disconnected(connections[c])
			}
		}
	}
	delete(connections, c)
}

func handleDisconnectNotification(c *rpc2.Client) {
	disconnected := c.DisconnectNotify()
	select {
	case <-disconnected:
		clearConnection(c)
	}
}

// Disconnect will close the OVSDB connection
func (ovs OvsdbClient) Disconnect() {
	ovs.rpcClient.Close()
}
<|endoftext|>


<commit_before>package conntrack

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"strconv"
	"syscall"
	"unsafe"
)

type nfgenmsg struct {
	Family  uint8  /* AF_xxx */
	Version uint8  /* nfnetlink version */
	ResID   uint16 /* resource id */
}

const (
	sizeofGenmsg = uint32(unsafe.Sizeof(nfgenmsg{})) // TODO
)

type ConntrackListReq struct {
	Header syscall.NlMsghdr
	Body   nfgenmsg
}

func (c *ConntrackListReq) toWireFormat() []byte {
	// adapted from syscall/NetlinkRouteRequest.toWireFormat
	b := make([]byte, c.Header.Len)
	*(*uint32)(unsafe.Pointer(&b[0:4][0])) = c.Header.Len
	*(*uint16)(unsafe.Pointer(&b[4:6][0])) = c.Header.Type
	*(*uint16)(unsafe.Pointer(&b[6:8][0])) = c.Header.Flags
	*(*uint32)(unsafe.Pointer(&b[8:12][0])) = c.Header.Seq
	*(*uint32)(unsafe.Pointer(&b[12:16][0])) = c.Header.Pid
	b[16] = byte(c.Body.Family)
	b[17] = byte(c.Body.Version)
	*(*uint16)(unsafe.Pointer(&b[18:20][0])) = c.Body.ResID
	return b
}

func connectNetfilter(groups uint32) (int, *syscall.SockaddrNetlink, error) {
	s, err := syscall.Socket(syscall.AF_NETLINK, syscall.SOCK_RAW, syscall.NETLINK_NETFILTER)
	if err != nil {
		return 0, nil, err
	}
	lsa := &syscall.SockaddrNetlink{
		Family: syscall.AF_NETLINK,
		Groups: groups,
	}
	if err := syscall.Bind(s, lsa); err != nil {
		return 0, nil, err
	}

	return s, lsa, nil
}

// Make syscall asking for all connections. Invoke 'cb' for each connection.
func queryAllConnections(cb func(Conn)) error {
	s, lsa, err := connectNetfilter(0)
	if err != nil {
		return err
	}
	defer syscall.Close(s)

	msg := ConntrackListReq{
		Header: syscall.NlMsghdr{
			Len:   syscall.NLMSG_HDRLEN + sizeofGenmsg,
			Type:  (NFNL_SUBSYS_CTNETLINK << 8) | uint16(IpctnlMsgCtGet),
			Flags: syscall.NLM_F_REQUEST | syscall.NLM_F_DUMP,
			Pid:   0,
			Seq:   0,
		},
		Body: nfgenmsg{
			Family:  syscall.AF_INET,
			Version: NFNETLINK_V0,
			ResID:   0,
		},
	}
	wb := msg.toWireFormat()
	// fmt.Printf("msg bytes: %q\n", wb)
	if err := syscall.Sendto(s, wb, 0, lsa); err != nil {
		return err
	}

	return readMsgs(s, cb)
}

// Lists all the connections that conntrack is tracking.
func Connections() ([]Conn, error) {
	var conns []Conn
	queryAllConnections(func(c Conn) {
		conns = append(conns, c)
	})
	return conns, nil
}

// Established lists all established TCP connections.
func Established() ([]ConnTCP, error) {
	var conns []ConnTCP
	local := localIPs()
	err := queryAllConnections(func(c Conn) {
		if c.MsgType != NfctMsgUpdate {
			fmt.Printf("msg isn't an update: %d\n", c.MsgType)
			return
		}
		if c.TCPState != "ESTABLISHED" {
			// fmt.Printf("state isn't ESTABLISHED: %s\n", c.TCPState)
			return
		}
		if tc := c.ConnTCP(local); tc != nil {
			conns = append(conns, *tc)
		}
	})
	if err != nil {
		return nil, err
	}
	return conns, nil
}

// Follow gives a channel with all changes.
func Follow() (<-chan Conn, func(), error) {
	s, _, err := connectNetfilter(
		NF_NETLINK_CONNTRACK_NEW | NF_NETLINK_CONNTRACK_UPDATE |
		NF_NETLINK_CONNTRACK_DESTROY)
	stop := func() {
		syscall.Close(s)
	}
	if err != nil {
		return nil, stop, err
	}

	res := make(chan Conn, 1)
	go func() {
		defer syscall.Close(s)
		if err := readMsgs(s, func(c Conn) {
			// if conn.TCPState != 3 {
			// // 3 is TCP established.
			// continue
			// }
			res <- c
		}); err != nil {
			panic(err)
		}
	}()
	return res, stop, nil
}

func readMsgs(s int, cb func(Conn)) error {
	rb := make([]byte, 2*syscall.Getpagesize())
	for {
		nr, _, err := syscall.Recvfrom(s, rb, 0)
		if err != nil {
			return err
		}

		msgs, err := syscall.ParseNetlinkMessage(rb[:nr])
		if err != nil {
			return err
		}
		for _, msg := range msgs {
			if err := nfnlIsError(msg.Header); err != nil {
				return fmt.Errorf("msg is some error: %s\n", err)
			}
			if nflnSubsysID(msg.Header.Type) != NFNL_SUBSYS_CTNETLINK {
				return fmt.Errorf(
					"unexpected subsys_id: %d\n",
					nflnSubsysID(msg.Header.Type),
				)
			}
			conn, err := parsePayload(msg.Data[sizeofGenmsg:])
			if err != nil {
				return err
			}

			// Taken from conntrack/parse.c:__parse_message_type
			switch CntlMsgTypes(nflnMsgType(msg.Header.Type)) {
			case IpctnlMsgCtNew:
				conn.MsgType = NfctMsgUpdate
				if msg.Header.Flags&(syscall.NLM_F_CREATE|syscall.NLM_F_EXCL) > 0 {
					conn.MsgType = NfctMsgNew
				}
			case IpctnlMsgCtDelete:
				conn.MsgType = NfctMsgDestroy
			}

			cb(*conn)
		}
	}
}

type Conn struct {
	MsgType  NfConntrackMsg
	Proto    int
	Src      net.IP
	SrcPort  uint16
	Dst      net.IP
	DstPort  uint16
	TCPState string

	// ICMP stuff.
	SrcIcmpId   uint16
	SrcIcmpType uint8
	SrcIcmpCode uint8
	DestIcmpId   uint16
	DestIcmpType uint8
	DestIcmpCode uint8

	// ct.mark, used to set permission type of the flow.
	CtMark uint32

	// ct.id, used to identify connections.
	CtId uint32

	// For multitenancy.
	Zone uint16

	// Flow stats.
	ReplyPktLen uint64
	ReplyPktCount uint64
	OrigPktLen uint64
	OrigPktCount uint64
}

// ConnTCP decides which way this connection is going and makes a ConnTCP.
func (c Conn) ConnTCP(local map[string]struct{}) *ConnTCP {
	// conntrack gives us all connections, even things passing through, but it
	// doesn't tell us what the local IP is. So we use `local` as a guide
	// what's local.
	src := c.Src.String()
	dst := c.Dst.String()
	_, srcLocal := local[src]
	_, dstLocal := local[dst]
	// If both are local we must just order things predictably.
	if srcLocal && dstLocal {
		srcLocal = c.SrcPort < c.DstPort
	}
	if srcLocal {
		return &ConnTCP{
			Local:      src,
			LocalPort:  strconv.Itoa(int(c.SrcPort)),
			Remote:     dst,
			RemotePort: strconv.Itoa(int(c.DstPort)),
		}
	}
	if dstLocal {
		return &ConnTCP{
			Local:      dst,
			LocalPort:  strconv.Itoa(int(c.DstPort)),
			Remote:     src,
			RemotePort: strconv.Itoa(int(c.SrcPort)),
		}
	}
	// Neither is local. conntrack also reports NAT connections.
	return nil
}

func parsePayload(b []byte) (*Conn, error) {
	// Most of this comes from libnetfilter_conntrack/src/conntrack/parse_mnl.c
	conn := &Conn{}
	attrs, err := parseAttrs(b)
	if err != nil {
		return conn, err
	}
	for _, attr := range attrs {
		switch CtattrType(attr.Typ) {
		case CtaTupleOrig:
			parseTuple(attr.Msg, conn, true /* isSrc */)
		case CtaTupleReply:
			parseTuple(attr.Msg, conn, false /* isSrc */)
		case CtaCountersOrig:
			conn.OrigPktLen, conn.OrigPktCount, _ = parseCounters(attr.Msg)
		case CtaCountersReply:
			conn.ReplyPktLen, conn.ReplyPktCount, _ = parseCounters(attr.Msg)
		case CtaStatus:
			// These are ip_conntrack_status
			// status := binary.BigEndian.Uint32(attr.Msg)
			// fmt.Printf("It's status %d\n", status)
		case CtaProtoinfo:
			parseProtoinfo(attr.Msg, conn)
		case CtaMark:
			conn.CtMark = binary.BigEndian.Uint32(attr.Msg)
		case CtaZone:
			conn.Zone = binary.BigEndian.Uint16(attr.Msg)
		case CtaId:
			conn.CtId = binary.BigEndian.Uint32(attr.Msg)
		}
	}
	return conn, nil
}

func parseTuple(b []byte, conn *Conn, isSrc bool) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}
	for _, attr := range attrs {
		// fmt.Printf("pl: %d, type: %d, multi: %t, bigend: %t\n", len(attr.Msg), attr.Typ, attr.IsNested, attr.IsNetByteorder)
		switch CtattrTuple(attr.Typ) {
		case CtaTupleUnspec:
			// fmt.Printf("It's a tuple unspec\n")
		case CtaTupleIp:
			// fmt.Printf("It's a tuple IP\n")
			if err := parseIP(attr.Msg, conn); err != nil {
				return err
			}
		case CtaTupleProto:
			// fmt.Printf("It's a tuple proto\n")
			parseProto(attr.Msg, conn, isSrc)
		}
	}
	return nil
}

func parseCounters(b []byte) (uint64, uint64, error) {
	attrs, err := parseAttrs(b)
	if err != nil {
		return 0, 0, fmt.Errorf("invalid tuple attr: %s", err)
	}
	packets := uint64(0)
	bytes := uint64(0)
	for _, attr := range attrs {
		switch CtattrCounters(attr.Typ) {
		case CtaCountersPackets:
			packets = binary.BigEndian.Uint64(attr.Msg)
		case CtaCountersBytes:
			bytes = binary.BigEndian.Uint64(attr.Msg)
		}
	}
	return packets, bytes, nil
}

func parseIP(b []byte, conn *Conn) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}
	for _, attr := range attrs {
		switch CtattrIp(attr.Typ) {
		case CtaIpV4Src:
			conn.Src = net.IP(attr.Msg) // TODO: copy so we can reuse the buffer?
		case CtaIpV4Dst:
			conn.Dst = net.IP(attr.Msg) // TODO: copy so we can reuse the buffer?
		case CtaIpV6Src:
			// TODO
		case CtaIpV6Dst:
			// TODO
		}
	}
	return nil
}

func parseProto(b []byte, conn *Conn, isSrc bool) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}

	for _, attr := range attrs {
		switch CtattrL4proto(attr.Typ) {
		// Protocol number.
		case CtaProtoNum:
			conn.Proto = int(uint8(attr.Msg[0]))

		// TCP stuff.
		case CtaProtoSrcPort:
			conn.SrcPort = binary.BigEndian.Uint16(attr.Msg)
		case CtaProtoDstPort:
			conn.DstPort = binary.BigEndian.Uint16(attr.Msg)

		// ICMP stuff.
		case CtaProtoIcmpId:
			if isSrc {
				conn.SrcIcmpId = binary.BigEndian.Uint16(attr.Msg)
			} else {
				conn.DestIcmpId = binary.BigEndian.Uint16(attr.Msg)
			}
		case CtaProtoIcmpType:
			bufreader := bytes.NewReader(attr.Msg)
			if isSrc {
				binary.Read(bufreader, binary.BigEndian, &conn.SrcIcmpType)
			} else {
				binary.Read(bufreader, binary.BigEndian, &conn.DestIcmpType)
			}
		case CtaProtoIcmpCode:
			bufreader := bytes.NewReader(attr.Msg)
			if isSrc {
				binary.Read(bufreader, binary.BigEndian, &conn.SrcIcmpCode)
			} else {
				binary.Read(bufreader, binary.BigEndian, &conn.DestIcmpCode)
			}
		}
	}

	return nil
}

func parseProtoinfo(b []byte, conn *Conn) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}
	for _, attr := range attrs {
		switch CtattrProtoinfo(attr.Typ) {
		case CtaProtoinfoTcp:
			if err := parseProtoinfoTCP(attr.Msg, conn); err != nil {
				return err
			}
		default:
			// we're not interested in other protocols
		}
	}
	return nil
}

func parseProtoinfoTCP(b []byte, conn *Conn) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}
	for _, attr := range attrs {
		switch CtattrProtoinfoTcp(attr.Typ) {
		case CtaProtoinfoTcpState:
			conn.TCPState = tcpState[uint8(attr.Msg[0])]
		default:
			// not interested
		}
	}
	return nil
}
<commit_msg>Return ENOBUF error to client<commit_after>package conntrack

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"strconv"
	"syscall"
	"unsafe"
)

type nfgenmsg struct {
	Family  uint8  /* AF_xxx */
	Version uint8  /* nfnetlink version */
	ResID   uint16 /* resource id */
}

const (
	sizeofGenmsg = uint32(unsafe.Sizeof(nfgenmsg{})) // TODO
)

type ConntrackListReq struct {
	Header syscall.NlMsghdr
	Body   nfgenmsg
}

func (c *ConntrackListReq) toWireFormat() []byte {
	// adapted from syscall/NetlinkRouteRequest.toWireFormat
	b := make([]byte, c.Header.Len)
	*(*uint32)(unsafe.Pointer(&b[0:4][0])) = c.Header.Len
	*(*uint16)(unsafe.Pointer(&b[4:6][0])) = c.Header.Type
	*(*uint16)(unsafe.Pointer(&b[6:8][0])) = c.Header.Flags
	*(*uint32)(unsafe.Pointer(&b[8:12][0])) = c.Header.Seq
	*(*uint32)(unsafe.Pointer(&b[12:16][0])) = c.Header.Pid
	b[16] = byte(c.Body.Family)
	b[17] = byte(c.Body.Version)
	*(*uint16)(unsafe.Pointer(&b[18:20][0])) = c.Body.ResID
	return b
}

func connectNetfilter(groups uint32) (int, *syscall.SockaddrNetlink, error) {
	s, err := syscall.Socket(syscall.AF_NETLINK, syscall.SOCK_RAW, syscall.NETLINK_NETFILTER)
	if err != nil {
		return 0, nil, err
	}
	lsa := &syscall.SockaddrNetlink{
		Family: syscall.AF_NETLINK,
		Groups: groups,
	}
	if err := syscall.Bind(s, lsa); err != nil {
		return 0, nil, err
	}

	return s, lsa, nil
}

// Make syscall asking for all connections. Invoke 'cb' for each connection.
func queryAllConnections(cb func(Conn)) error {
	s, lsa, err := connectNetfilter(0)
	if err != nil {
		return err
	}
	defer syscall.Close(s)

	msg := ConntrackListReq{
		Header: syscall.NlMsghdr{
			Len:   syscall.NLMSG_HDRLEN + sizeofGenmsg,
			Type:  (NFNL_SUBSYS_CTNETLINK << 8) | uint16(IpctnlMsgCtGet),
			Flags: syscall.NLM_F_REQUEST | syscall.NLM_F_DUMP,
			Pid:   0,
			Seq:   0,
		},
		Body: nfgenmsg{
			Family:  syscall.AF_INET,
			Version: NFNETLINK_V0,
			ResID:   0,
		},
	}
	wb := msg.toWireFormat()
	// fmt.Printf("msg bytes: %q\n", wb)
	if err := syscall.Sendto(s, wb, 0, lsa); err != nil {
		return err
	}

	return readMsgs(s, cb)
}

// Lists all the connections that conntrack is tracking.
func Connections() ([]Conn, error) {
	var conns []Conn
	queryAllConnections(func(c Conn) {
		conns = append(conns, c)
	})
	return conns, nil
}

// Established lists all established TCP connections.
func Established() ([]ConnTCP, error) {
	var conns []ConnTCP
	local := localIPs()
	err := queryAllConnections(func(c Conn) {
		if c.MsgType != NfctMsgUpdate {
			fmt.Printf("msg isn't an update: %d\n", c.MsgType)
			return
		}
		if c.TCPState != "ESTABLISHED" {
			// fmt.Printf("state isn't ESTABLISHED: %s\n", c.TCPState)
			return
		}
		if tc := c.ConnTCP(local); tc != nil {
			conns = append(conns, *tc)
		}
	})
	if err != nil {
		return nil, err
	}
	return conns, nil
}

// Follow gives a channel with all changes.
func Follow() (<-chan Conn, func(), error) {
	s, _, err := connectNetfilter(
		NF_NETLINK_CONNTRACK_NEW | NF_NETLINK_CONNTRACK_UPDATE |
		NF_NETLINK_CONNTRACK_DESTROY)
	stop := func() {
		syscall.Close(s)
	}
	if err != nil {
		return nil, stop, err
	}

	res := make(chan Conn, 1)
	go func() {
		defer syscall.Close(s)
		if err := readMsgs(s, func(c Conn) {
			// if conn.TCPState != 3 {
			// // 3 is TCP established.
			// continue
			// }
			res <- c
		}); err != nil {
			panic(err)
		}
	}()
	return res, stop, nil
}

func readMsgs(s int, cb func(Conn)) error {
	rb := make([]byte, 2*syscall.Getpagesize())
	for {
		nr, _, err := syscall.Recvfrom(s, rb, 0)
		if err == syscall.ENOBUFS {
			// ENOBUF means we miss some events here. No way around it. That's life.
			cb(Conn{Err: syscall.ENOBUFS})
			continue
		} else if err != nil {
			return err
		}

		msgs, err := syscall.ParseNetlinkMessage(rb[:nr])
		if err != nil {
			return err
		}
		for _, msg := range msgs {
			if err := nfnlIsError(msg.Header); err != nil {
				return fmt.Errorf("msg is some error: %s\n", err)
			}
			if nflnSubsysID(msg.Header.Type) != NFNL_SUBSYS_CTNETLINK {
				return fmt.Errorf(
					"unexpected subsys_id: %d\n",
					nflnSubsysID(msg.Header.Type),
				)
			}
			conn, err := parsePayload(msg.Data[sizeofGenmsg:])
			if err != nil {
				return err
			}

			// Taken from conntrack/parse.c:__parse_message_type
			switch CntlMsgTypes(nflnMsgType(msg.Header.Type)) {
			case IpctnlMsgCtNew:
				conn.MsgType = NfctMsgUpdate
				if msg.Header.Flags&(syscall.NLM_F_CREATE|syscall.NLM_F_EXCL) > 0 {
					conn.MsgType = NfctMsgNew
				}
			case IpctnlMsgCtDelete:
				conn.MsgType = NfctMsgDestroy
			}

			cb(*conn)
		}
	}
}

type Conn struct {
	MsgType  NfConntrackMsg
	Proto    int
	Src      net.IP
	SrcPort  uint16
	Dst      net.IP
	DstPort  uint16
	TCPState string

	// ICMP stuff.
	SrcIcmpId   uint16
	SrcIcmpType uint8
	SrcIcmpCode uint8
	DestIcmpId   uint16
	DestIcmpType uint8
	DestIcmpCode uint8

	// ct.mark, used to set permission type of the flow.
	CtMark uint32

	// ct.id, used to identify connections.
	CtId uint32

	// For multitenancy.
	Zone uint16

	// Flow stats.
	ReplyPktLen uint64
	ReplyPktCount uint64
	OrigPktLen uint64
	OrigPktCount uint64

	// Error, if any.
	Err error
}

// ConnTCP decides which way this connection is going and makes a ConnTCP.
func (c Conn) ConnTCP(local map[string]struct{}) *ConnTCP {
	// conntrack gives us all connections, even things passing through, but it
	// doesn't tell us what the local IP is. So we use `local` as a guide
	// what's local.
	src := c.Src.String()
	dst := c.Dst.String()
	_, srcLocal := local[src]
	_, dstLocal := local[dst]
	// If both are local we must just order things predictably.
	if srcLocal && dstLocal {
		srcLocal = c.SrcPort < c.DstPort
	}
	if srcLocal {
		return &ConnTCP{
			Local:      src,
			LocalPort:  strconv.Itoa(int(c.SrcPort)),
			Remote:     dst,
			RemotePort: strconv.Itoa(int(c.DstPort)),
		}
	}
	if dstLocal {
		return &ConnTCP{
			Local:      dst,
			LocalPort:  strconv.Itoa(int(c.DstPort)),
			Remote:     src,
			RemotePort: strconv.Itoa(int(c.SrcPort)),
		}
	}
	// Neither is local. conntrack also reports NAT connections.
	return nil
}

func parsePayload(b []byte) (*Conn, error) {
	// Most of this comes from libnetfilter_conntrack/src/conntrack/parse_mnl.c
	conn := &Conn{}
	attrs, err := parseAttrs(b)
	if err != nil {
		return conn, err
	}
	for _, attr := range attrs {
		switch CtattrType(attr.Typ) {
		case CtaTupleOrig:
			parseTuple(attr.Msg, conn, true /* isSrc */)
		case CtaTupleReply:
			parseTuple(attr.Msg, conn, false /* isSrc */)
		case CtaCountersOrig:
			conn.OrigPktLen, conn.OrigPktCount, _ = parseCounters(attr.Msg)
		case CtaCountersReply:
			conn.ReplyPktLen, conn.ReplyPktCount, _ = parseCounters(attr.Msg)
		case CtaStatus:
			// These are ip_conntrack_status
			// status := binary.BigEndian.Uint32(attr.Msg)
			// fmt.Printf("It's status %d\n", status)
		case CtaProtoinfo:
			parseProtoinfo(attr.Msg, conn)
		case CtaMark:
			conn.CtMark = binary.BigEndian.Uint32(attr.Msg)
		case CtaZone:
			conn.Zone = binary.BigEndian.Uint16(attr.Msg)
		case CtaId:
			conn.CtId = binary.BigEndian.Uint32(attr.Msg)
		}
	}
	return conn, nil
}

func parseTuple(b []byte, conn *Conn, isSrc bool) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}
	for _, attr := range attrs {
		// fmt.Printf("pl: %d, type: %d, multi: %t, bigend: %t\n", len(attr.Msg), attr.Typ, attr.IsNested, attr.IsNetByteorder)
		switch CtattrTuple(attr.Typ) {
		case CtaTupleUnspec:
			// fmt.Printf("It's a tuple unspec\n")
		case CtaTupleIp:
			// fmt.Printf("It's a tuple IP\n")
			if err := parseIP(attr.Msg, conn); err != nil {
				return err
			}
		case CtaTupleProto:
			// fmt.Printf("It's a tuple proto\n")
			parseProto(attr.Msg, conn, isSrc)
		}
	}
	return nil
}

func parseCounters(b []byte) (uint64, uint64, error) {
	attrs, err := parseAttrs(b)
	if err != nil {
		return 0, 0, fmt.Errorf("invalid tuple attr: %s", err)
	}
	packets := uint64(0)
	bytes := uint64(0)
	for _, attr := range attrs {
		switch CtattrCounters(attr.Typ) {
		case CtaCountersPackets:
			packets = binary.BigEndian.Uint64(attr.Msg)
		case CtaCountersBytes:
			bytes = binary.BigEndian.Uint64(attr.Msg)
		}
	}
	return packets, bytes, nil
}

func parseIP(b []byte, conn *Conn) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}
	for _, attr := range attrs {
		switch CtattrIp(attr.Typ) {
		case CtaIpV4Src:
			conn.Src = net.IP(attr.Msg) // TODO: copy so we can reuse the buffer?
		case CtaIpV4Dst:
			conn.Dst = net.IP(attr.Msg) // TODO: copy so we can reuse the buffer?
		case CtaIpV6Src:
			// TODO
		case CtaIpV6Dst:
			// TODO
		}
	}
	return nil
}

func parseProto(b []byte, conn *Conn, isSrc bool) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}

	for _, attr := range attrs {
		switch CtattrL4proto(attr.Typ) {
		// Protocol number.
		case CtaProtoNum:
			conn.Proto = int(uint8(attr.Msg[0]))

		// TCP stuff.
		case CtaProtoSrcPort:
			conn.SrcPort = binary.BigEndian.Uint16(attr.Msg)
		case CtaProtoDstPort:
			conn.DstPort = binary.BigEndian.Uint16(attr.Msg)

		// ICMP stuff.
		case CtaProtoIcmpId:
			if isSrc {
				conn.SrcIcmpId = binary.BigEndian.Uint16(attr.Msg)
			} else {
				conn.DestIcmpId = binary.BigEndian.Uint16(attr.Msg)
			}
		case CtaProtoIcmpType:
			bufreader := bytes.NewReader(attr.Msg)
			if isSrc {
				binary.Read(bufreader, binary.BigEndian, &conn.SrcIcmpType)
			} else {
				binary.Read(bufreader, binary.BigEndian, &conn.DestIcmpType)
			}
		case CtaProtoIcmpCode:
			bufreader := bytes.NewReader(attr.Msg)
			if isSrc {
				binary.Read(bufreader, binary.BigEndian, &conn.SrcIcmpCode)
			} else {
				binary.Read(bufreader, binary.BigEndian, &conn.DestIcmpCode)
			}
		}
	}

	return nil
}

func parseProtoinfo(b []byte, conn *Conn) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}
	for _, attr := range attrs {
		switch CtattrProtoinfo(attr.Typ) {
		case CtaProtoinfoTcp:
			if err := parseProtoinfoTCP(attr.Msg, conn); err != nil {
				return err
			}
		default:
			// we're not interested in other protocols
		}
	}
	return nil
}

func parseProtoinfoTCP(b []byte, conn *Conn) error {
	attrs, err := parseAttrs(b)
	if err != nil {
		return fmt.Errorf("invalid tuple attr: %s", err)
	}
	for _, attr := range attrs {
		switch CtattrProtoinfoTcp(attr.Typ) {
		case CtaProtoinfoTcpState:
			conn.TCPState = tcpState[uint8(attr.Msg[0])]
		default:
			// not interested
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package dhcp4client

import (
	"bytes"
	"crypto/rand"
	"net"
	"time"

	"github.com/d2g/dhcp4"
)

const (
	MaxDHCPLen = 576
)

type Client struct {
	hardwareAddr  net.HardwareAddr //The HardwareAddr to send in the request.
	ignoreServers []net.IP         //List of Servers to Ignore requests from.
	timeout       time.Duration    //Time before we timeout.
	broadcast     bool             //Set the Bcast flag in BOOTP Flags
	connection    connection       //The Connection Method to use
}

/*
 * Abstracts the type of underlying socket used
 */
type connection interface {
	Close() error
	Write(packet []byte) error
	ReadFrom() ([]byte, net.IP, error)
	SetReadTimeout(t time.Duration) error
}

func New(options ...func(*Client) error) (*Client, error) {
	c := Client{
		timeout:   time.Second * 10,
		broadcast: true,
	}

	err := c.SetOption(options...)
	if err != nil {
		return nil, err
	}

	//if connection hasn't been set as an option create the default.
	if c.connection == nil {
		conn, err := NewInetSock()
		if err != nil {
			return nil, err
		}
		c.connection = conn
	}

	return &c, nil
}

func (c *Client) SetOption(options ...func(*Client) error) error {
	for _, opt := range options {
		if err := opt(c); err != nil {
			return err
		}
	}
	return nil
}

func Timeout(t time.Duration) func(*Client) error {
	return func(c *Client) error {
		c.timeout = t
		return nil
	}
}

func IgnoreServers(s []net.IP) func(*Client) error {
	return func(c *Client) error {
		c.ignoreServers = s
		return nil
	}
}

func HardwareAddr(h net.HardwareAddr) func(*Client) error {
	return func(c *Client) error {
		c.hardwareAddr = h
		return nil
	}
}

func Broadcast(b bool) func(*Client) error {
	return func(c *Client) error {
		c.broadcast = b
		return nil
	}
}

func Connection(conn connection) func(*Client) error {
	return func(c *Client) error {
		c.connection = conn
		return nil
	}
}

/*
 * Close Connections
 */
func (c *Client) Close() error {
	if c.connection != nil {
		return c.connection.Close()
	}
	return nil
}

/*
 * Send the Discovery Packet to the Broadcast Channel
 */
func (c *Client) SendDiscoverPacket() (dhcp4.Packet, error) {
	discoveryPacket := c.DiscoverPacket()
	discoveryPacket.PadToMinSize()

	return discoveryPacket, c.SendPacket(discoveryPacket)
}

/*
 * Retreive Offer...
 * Wait for the offer for a specific Discovery Packet.
 */
func (c *Client) GetOffer(discoverPacket *dhcp4.Packet) (dhcp4.Packet, error) {
	for {
		c.connection.SetReadTimeout(c.timeout)
		readBuffer, source, err := c.connection.ReadFrom()
		if err != nil {
			return dhcp4.Packet{}, err
		}

		offerPacket := dhcp4.Packet(readBuffer)
		offerPacketOptions := offerPacket.ParseOptions()

		// Ignore Servers in my Ignore list
		for _, ignoreServer := range c.ignoreServers {
			if source.Equal(ignoreServer) {
				continue
			}

			if offerPacket.SIAddr().Equal(ignoreServer) {
				continue
			}
		}

		if len(offerPacketOptions[dhcp4.OptionDHCPMessageType]) < 1 || dhcp4.MessageType(offerPacketOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.Offer || !bytes.Equal(discoverPacket.XId(), offerPacket.XId()) {
			continue
		}

		return offerPacket, nil
	}

}

/*
 * Send Request Based On the offer Received.
 */
func (c *Client) SendRequest(offerPacket *dhcp4.Packet) (dhcp4.Packet, error) {
	requestPacket := c.RequestPacket(offerPacket)
	requestPacket.PadToMinSize()

	return requestPacket, c.SendPacket(requestPacket)
}

/*
 * Retreive Acknowledgement
 * Wait for the offer for a specific Request Packet.
 */
func (c *Client) GetAcknowledgement(requestPacket *dhcp4.Packet) (dhcp4.Packet, error) {
	for {
		c.connection.SetReadTimeout(c.timeout)
		readBuffer, source, err := c.connection.ReadFrom()
		if err != nil {
			return dhcp4.Packet{}, err
		}

		acknowledgementPacket := dhcp4.Packet(readBuffer)
		acknowledgementPacketOptions := acknowledgementPacket.ParseOptions()

		// Ignore Servers in my Ignore list
		for _, ignoreServer := range c.ignoreServers {
			if source.Equal(ignoreServer) {
				continue
			}

			if acknowledgementPacket.SIAddr().Equal(ignoreServer) {
				continue
			}
		}

		if !bytes.Equal(requestPacket.XId(), acknowledgementPacket.XId()) || len(acknowledgementPacketOptions[dhcp4.OptionDHCPMessageType]) < 1 || (dhcp4.MessageType(acknowledgementPacketOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.ACK && dhcp4.MessageType(acknowledgementPacketOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.NAK) {
			continue
		}

		return acknowledgementPacket, nil
	}
}

/*
 * Send a DHCP Packet.
 */
func (c *Client) SendPacket(packet dhcp4.Packet) error {
	return c.connection.Write(packet)
}

/*
 * Create Discover Packet
 */
func (c *Client) DiscoverPacket() dhcp4.Packet {
	messageid := make([]byte, 4)
	if _, err := rand.Read(messageid); err != nil {
		panic(err)
	}

	packet := dhcp4.NewPacket(dhcp4.BootRequest)
	packet.SetCHAddr(c.hardwareAddr)
	packet.SetXId(messageid)
	packet.SetBroadcast(c.broadcast)

	packet.AddOption(dhcp4.OptionDHCPMessageType, []byte{byte(dhcp4.Discover)})
	//packet.PadToMinSize()
	return packet
}

/*
 * Create Request Packet
 */
func (c *Client) RequestPacket(offerPacket *dhcp4.Packet) dhcp4.Packet {
	offerOptions := offerPacket.ParseOptions()

	packet := dhcp4.NewPacket(dhcp4.BootRequest)
	packet.SetCHAddr(c.hardwareAddr)

	packet.SetXId(offerPacket.XId())
	packet.SetCIAddr(offerPacket.CIAddr())
	packet.SetSIAddr(offerPacket.SIAddr())

	packet.SetBroadcast(c.broadcast)
	packet.AddOption(dhcp4.OptionDHCPMessageType, []byte{byte(dhcp4.Request)})
	packet.AddOption(dhcp4.OptionRequestedIPAddress, (offerPacket.YIAddr()).To4())
	packet.AddOption(dhcp4.OptionServerIdentifier, offerOptions[dhcp4.OptionServerIdentifier])

	//packet.PadToMinSize()
	return packet
}

/*
 * Create Request Packet For a Renew
 */
func (c *Client) RenewalRequestPacket(acknowledgement *dhcp4.Packet) dhcp4.Packet {
	messageid := make([]byte, 4)
	if _, err := rand.Read(messageid); err != nil {
		panic(err)
	}

	acknowledgementOptions := acknowledgement.ParseOptions()

	packet := dhcp4.NewPacket(dhcp4.BootRequest)
	packet.SetCHAddr(acknowledgement.CHAddr())

	packet.SetXId(messageid)
	packet.SetCIAddr(acknowledgement.YIAddr())
	packet.SetSIAddr(acknowledgement.SIAddr())

	packet.SetBroadcast(c.broadcast)
	packet.AddOption(dhcp4.OptionDHCPMessageType, []byte{byte(dhcp4.Request)})
	packet.AddOption(dhcp4.OptionRequestedIPAddress, (acknowledgement.YIAddr()).To4())
	packet.AddOption(dhcp4.OptionServerIdentifier, acknowledgementOptions[dhcp4.OptionServerIdentifier])

	//packet.PadToMinSize()
	return packet
}

/*
 * Lets do a Full DHCP Request.
 */
func (c *Client) Request() (bool, dhcp4.Packet, error) {
	discoveryPacket, err := c.SendDiscoverPacket()
	if err != nil {
		return false, discoveryPacket, err
	}

	offerPacket, err := c.GetOffer(&discoveryPacket)
	if err != nil {
		return false, offerPacket, err
	}

	requestPacket, err := c.SendRequest(&offerPacket)
	if err != nil {
		return false, requestPacket, err
	}

	acknowledgement, err := c.GetAcknowledgement(&requestPacket)
	if err != nil {
		return false, acknowledgement, err
	}

	acknowledgementOptions := acknowledgement.ParseOptions()
	if dhcp4.MessageType(acknowledgementOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.ACK {
		return false, acknowledgement, nil
	}

	return true, acknowledgement, nil
}

/*
 * Renew a lease backed on the Acknowledgement Packet.
 * Returns Sucessfull, The AcknoledgementPacket, Any Errors
 */
func (c *Client) Renew(acknowledgement dhcp4.Packet) (bool, dhcp4.Packet, error) {
	renewRequest := c.RenewalRequestPacket(&acknowledgement)
	renewRequest.PadToMinSize()

	err := c.SendPacket(renewRequest)
	if err != nil {
		return false, renewRequest, err
	}

	newAcknowledgement, err := c.GetAcknowledgement(&acknowledgement)
	if err != nil {
		return false, newAcknowledgement, err
	}

	newAcknowledgementOptions := newAcknowledgement.ParseOptions()
	if dhcp4.MessageType(newAcknowledgementOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.ACK {
		return false, newAcknowledgement, nil
	}

	return true, newAcknowledgement, nil
}
<commit_msg>add DHCPRELEASE support<commit_after>package dhcp4client

import (
	"bytes"
	"crypto/rand"
	"net"
	"time"

	"github.com/d2g/dhcp4"
)

const (
	MaxDHCPLen = 576
)

type Client struct {
	hardwareAddr  net.HardwareAddr //The HardwareAddr to send in the request.
	ignoreServers []net.IP         //List of Servers to Ignore requests from.
	timeout       time.Duration    //Time before we timeout.
	broadcast     bool             //Set the Bcast flag in BOOTP Flags
	connection    connection       //The Connection Method to use
}

/*
 * Abstracts the type of underlying socket used
 */
type connection interface {
	Close() error
	Write(packet []byte) error
	ReadFrom() ([]byte, net.IP, error)
	SetReadTimeout(t time.Duration) error
}

func New(options ...func(*Client) error) (*Client, error) {
	c := Client{
		timeout:   time.Second * 10,
		broadcast: true,
	}

	err := c.SetOption(options...)
	if err != nil {
		return nil, err
	}

	//if connection hasn't been set as an option create the default.
	if c.connection == nil {
		conn, err := NewInetSock()
		if err != nil {
			return nil, err
		}
		c.connection = conn
	}

	return &c, nil
}

func (c *Client) SetOption(options ...func(*Client) error) error {
	for _, opt := range options {
		if err := opt(c); err != nil {
			return err
		}
	}
	return nil
}

func Timeout(t time.Duration) func(*Client) error {
	return func(c *Client) error {
		c.timeout = t
		return nil
	}
}

func IgnoreServers(s []net.IP) func(*Client) error {
	return func(c *Client) error {
		c.ignoreServers = s
		return nil
	}
}

func HardwareAddr(h net.HardwareAddr) func(*Client) error {
	return func(c *Client) error {
		c.hardwareAddr = h
		return nil
	}
}

func Broadcast(b bool) func(*Client) error {
	return func(c *Client) error {
		c.broadcast = b
		return nil
	}
}

func Connection(conn connection) func(*Client) error {
	return func(c *Client) error {
		c.connection = conn
		return nil
	}
}

/*
 * Close Connections
 */
func (c *Client) Close() error {
	if c.connection != nil {
		return c.connection.Close()
	}
	return nil
}

/*
 * Send the Discovery Packet to the Broadcast Channel
 */
func (c *Client) SendDiscoverPacket() (dhcp4.Packet, error) {
	discoveryPacket := c.DiscoverPacket()
	discoveryPacket.PadToMinSize()

	return discoveryPacket, c.SendPacket(discoveryPacket)
}

/*
 * Retreive Offer...
 * Wait for the offer for a specific Discovery Packet.
 */
func (c *Client) GetOffer(discoverPacket *dhcp4.Packet) (dhcp4.Packet, error) {
	for {
		c.connection.SetReadTimeout(c.timeout)
		readBuffer, source, err := c.connection.ReadFrom()
		if err != nil {
			return dhcp4.Packet{}, err
		}

		offerPacket := dhcp4.Packet(readBuffer)
		offerPacketOptions := offerPacket.ParseOptions()

		// Ignore Servers in my Ignore list
		for _, ignoreServer := range c.ignoreServers {
			if source.Equal(ignoreServer) {
				continue
			}

			if offerPacket.SIAddr().Equal(ignoreServer) {
				continue
			}
		}

		if len(offerPacketOptions[dhcp4.OptionDHCPMessageType]) < 1 || dhcp4.MessageType(offerPacketOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.Offer || !bytes.Equal(discoverPacket.XId(), offerPacket.XId()) {
			continue
		}

		return offerPacket, nil
	}

}

/*
 * Send Request Based On the offer Received.
 */
func (c *Client) SendRequest(offerPacket *dhcp4.Packet) (dhcp4.Packet, error) {
	requestPacket := c.RequestPacket(offerPacket)
	requestPacket.PadToMinSize()

	return requestPacket, c.SendPacket(requestPacket)
}

/*
 * Retreive Acknowledgement
 * Wait for the offer for a specific Request Packet.
 */
func (c *Client) GetAcknowledgement(requestPacket *dhcp4.Packet) (dhcp4.Packet, error) {
	for {
		c.connection.SetReadTimeout(c.timeout)
		readBuffer, source, err := c.connection.ReadFrom()
		if err != nil {
			return dhcp4.Packet{}, err
		}

		acknowledgementPacket := dhcp4.Packet(readBuffer)
		acknowledgementPacketOptions := acknowledgementPacket.ParseOptions()

		// Ignore Servers in my Ignore list
		for _, ignoreServer := range c.ignoreServers {
			if source.Equal(ignoreServer) {
				continue
			}

			if acknowledgementPacket.SIAddr().Equal(ignoreServer) {
				continue
			}
		}

		if !bytes.Equal(requestPacket.XId(), acknowledgementPacket.XId()) || len(acknowledgementPacketOptions[dhcp4.OptionDHCPMessageType]) < 1 || (dhcp4.MessageType(acknowledgementPacketOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.ACK && dhcp4.MessageType(acknowledgementPacketOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.NAK) {
			continue
		}

		return acknowledgementPacket, nil
	}
}

/*
 * Send a DHCP Packet.
 */
func (c *Client) SendPacket(packet dhcp4.Packet) error {
	return c.connection.Write(packet)
}

/*
 * Create Discover Packet
 */
func (c *Client) DiscoverPacket() dhcp4.Packet {
	messageid := make([]byte, 4)
	if _, err := rand.Read(messageid); err != nil {
		panic(err)
	}

	packet := dhcp4.NewPacket(dhcp4.BootRequest)
	packet.SetCHAddr(c.hardwareAddr)
	packet.SetXId(messageid)
	packet.SetBroadcast(c.broadcast)

	packet.AddOption(dhcp4.OptionDHCPMessageType, []byte{byte(dhcp4.Discover)})
	//packet.PadToMinSize()
	return packet
}

/*
 * Create Request Packet
 */
func (c *Client) RequestPacket(offerPacket *dhcp4.Packet) dhcp4.Packet {
	offerOptions := offerPacket.ParseOptions()

	packet := dhcp4.NewPacket(dhcp4.BootRequest)
	packet.SetCHAddr(c.hardwareAddr)

	packet.SetXId(offerPacket.XId())
	packet.SetCIAddr(offerPacket.CIAddr())
	packet.SetSIAddr(offerPacket.SIAddr())

	packet.SetBroadcast(c.broadcast)
	packet.AddOption(dhcp4.OptionDHCPMessageType, []byte{byte(dhcp4.Request)})
	packet.AddOption(dhcp4.OptionRequestedIPAddress, (offerPacket.YIAddr()).To4())
	packet.AddOption(dhcp4.OptionServerIdentifier, offerOptions[dhcp4.OptionServerIdentifier])

	//packet.PadToMinSize()
	return packet
}

/*
 * Create Request Packet For a Renew
 */
func (c *Client) RenewalRequestPacket(acknowledgement *dhcp4.Packet) dhcp4.Packet {
	messageid := make([]byte, 4)
	if _, err := rand.Read(messageid); err != nil {
		panic(err)
	}

	acknowledgementOptions := acknowledgement.ParseOptions()

	packet := dhcp4.NewPacket(dhcp4.BootRequest)
	packet.SetCHAddr(acknowledgement.CHAddr())

	packet.SetXId(messageid)
	packet.SetCIAddr(acknowledgement.YIAddr())
	packet.SetSIAddr(acknowledgement.SIAddr())

	packet.SetBroadcast(c.broadcast)
	packet.AddOption(dhcp4.OptionDHCPMessageType, []byte{byte(dhcp4.Request)})
	packet.AddOption(dhcp4.OptionRequestedIPAddress, (acknowledgement.YIAddr()).To4())
	packet.AddOption(dhcp4.OptionServerIdentifier, acknowledgementOptions[dhcp4.OptionServerIdentifier])

	//packet.PadToMinSize()
	return packet
}

/*
 * Create Release Packet For a Release
 */
func (c *Client) ReleasePacket(acknowledgement *dhcp4.Packet) dhcp4.Packet {
	messageid := make([]byte, 4)
	if _, err := rand.Read(messageid); err != nil {
		panic(err)
	}

	acknowledgementOptions := acknowledgement.ParseOptions()

	packet := dhcp4.NewPacket(dhcp4.BootRequest)
	packet.SetCHAddr(acknowledgement.CHAddr())

	packet.SetXId(messageid)
	packet.SetCIAddr(acknowledgement.YIAddr())

	packet.AddOption(dhcp4.OptionDHCPMessageType, []byte{byte(dhcp4.Release)})
	packet.AddOption(dhcp4.OptionServerIdentifier, acknowledgementOptions[dhcp4.OptionServerIdentifier])

	//packet.PadToMinSize()
	return packet
}

/*
 * Lets do a Full DHCP Request.
 */
func (c *Client) Request() (bool, dhcp4.Packet, error) {
	discoveryPacket, err := c.SendDiscoverPacket()
	if err != nil {
		return false, discoveryPacket, err
	}

	offerPacket, err := c.GetOffer(&discoveryPacket)
	if err != nil {
		return false, offerPacket, err
	}

	requestPacket, err := c.SendRequest(&offerPacket)
	if err != nil {
		return false, requestPacket, err
	}

	acknowledgement, err := c.GetAcknowledgement(&requestPacket)
	if err != nil {
		return false, acknowledgement, err
	}

	acknowledgementOptions := acknowledgement.ParseOptions()
	if dhcp4.MessageType(acknowledgementOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.ACK {
		return false, acknowledgement, nil
	}

	return true, acknowledgement, nil
}

/*
 * Renew a lease backed on the Acknowledgement Packet.
 * Returns Sucessfull, The AcknoledgementPacket, Any Errors
 */
func (c *Client) Renew(acknowledgement dhcp4.Packet) (bool, dhcp4.Packet, error) {
	renewRequest := c.RenewalRequestPacket(&acknowledgement)
	renewRequest.PadToMinSize()

	err := c.SendPacket(renewRequest)
	if err != nil {
		return false, renewRequest, err
	}

	newAcknowledgement, err := c.GetAcknowledgement(&acknowledgement)
	if err != nil {
		return false, newAcknowledgement, err
	}

	newAcknowledgementOptions := newAcknowledgement.ParseOptions()
	if dhcp4.MessageType(newAcknowledgementOptions[dhcp4.OptionDHCPMessageType][0]) != dhcp4.ACK {
		return false, newAcknowledgement, nil
	}

	return true, newAcknowledgement, nil
}

/*
 * Release a lease backed on the Acknowledgement Packet.
 * Returns Any Errors
 */
func (c *Client) Release(acknowledgement dhcp4.Packet) error {
	release := c.ReleasePacket(&acknowledgement)
	release.PadToMinSize()

	return c.SendPacket(release)
}
<|endoftext|>


<commit_before>package ionic

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const maxPagingLimit = 100

// IonClient represnets a communication layer with the Ion Channel API
type IonClient struct {
	baseURL     *url.URL
	bearerToken string
	client      *http.Client
}

// IonResponse represents the response structure expected back from the Ion
// Channel API calls
type IonResponse struct {
	Data json.RawMessage `json:"data"`
	Meta Meta            `json:"meta"`
}

// Meta represents the metadata section of an IonResponse
type Meta struct {
	Version    string    `json:"version"`
	LastUpdate time.Time `json:"last_update"`
	TotalCount int       `json:"total_count"`
	Limit      int       `json:"limit"`
	Offset     int       `json:"offset"`
}

// Pagination represents the necessary elements for a paginated request
type Pagination struct {
	Offset int
	Limit  int
}

// AllItems is a convenience for requesting all items of a given entity
var AllItems = &Pagination{Offset: 0, Limit: -1}

// New takes the credentials required to talk to the API, and the base URL of
// the API and returns a client for talking to the API and an error if any
// issues instantiating the client are encountered
func New(secret string, baseURL string) (*IonClient, error) {
	u, err := url.Parse(baseURL)
	if err != nil {
		return nil, fmt.Errorf("cannot instantiate new ion client: %v", err.Error())
	}

	ic := &IonClient{
		baseURL:     u,
		bearerToken: secret,
		client:      &http.Client{},
	}

	return ic, nil
}

func (ic *IonClient) createURL(endpoint string, params *url.Values, pagination *Pagination) *url.URL {
	u := *ic.baseURL
	u.Path = endpoint

	vals := &url.Values{}
	if params != nil {
		vals = params
	}

	if pagination != nil {
		vals.Set("offset", strconv.Itoa(pagination.Offset))
		vals.Set("limit", strconv.Itoa(pagination.Limit))
	}

	u.RawQuery = vals.Encode()
	return &u
}

func (ic *IonClient) do(method, endpoint string, params *url.Values, payload []byte, pagination *Pagination) (json.RawMessage, error) {
	if pagination == nil || pagination.Limit > 0 {
		ir, err := ic._do(method, endpoint, params, payload, pagination)
		if err != nil {
			return nil, err
		}

		return ir.Data, nil
	}

	page := &Pagination{Limit: maxPagingLimit, Offset: 0}
	var data json.RawMessage
	data = append(data, []byte("[")...)

	total := 1
	for page.Offset < total {
		ir, err := ic._do(method, endpoint, params, payload, page)
		if err != nil {
			return nil, fmt.Errorf("trouble paging from API: %v", err.Error())
		}
		data = append(data, ir.Data[1:len(ir.Data)-1]...)
		data = append(data, []byte(",")...)
		page.Offset += maxPagingLimit
		total = ir.Meta.TotalCount
	}

	data = append(data[:len(data)-1], []byte("]")...)
	return data, nil
}

func (ic *IonClient) _do(method, endpoint string, params *url.Values, payload []byte, pagination *Pagination) (*IonResponse, error) {
	u := ic.createURL(endpoint, params, pagination)

	req, err := http.NewRequest(strings.ToUpper(method), u.String(), nil)

	if ic.bearerToken != "" {
		req.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ic.bearerToken))
	}

	resp, err := ic.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed http request: %v", err.Error())
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("bad response from API: %v", resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response body: %v", err.Error())
	}

	var ir IonResponse
	err = json.Unmarshal(body, &ir)
	if err != nil {
		return nil, fmt.Errorf("malformed response: %v", err.Error())
	}

	return &ir, nil
}

func (ic *IonClient) delete(endpoint string, params *url.Values) (json.RawMessage, error) {
	return ic.do("DELETE", endpoint, params, nil, nil)
}

func (ic *IonClient) get(endpoint string, params *url.Values, pagination *Pagination) (json.RawMessage, error) {
	return ic.do("GET", endpoint, params, nil, pagination)
}

func (ic *IonClient) post(endpoint string, params *url.Values, payload []byte) (json.RawMessage, error) {
	return ic.do("POST", endpoint, params, payload, nil)
}

func (ic *IonClient) put(endpoint string, params *url.Values, payload []byte) (json.RawMessage, error) {
	return ic.do("PUT", endpoint, params, payload, nil)
}
<commit_msg>allow passing in headers<commit_after>package ionic

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const maxPagingLimit = 100

// IonClient represnets a communication layer with the Ion Channel API
type IonClient struct {
	baseURL     *url.URL
	bearerToken string
	client      *http.Client
}

// IonResponse represents the response structure expected back from the Ion
// Channel API calls
type IonResponse struct {
	Data json.RawMessage `json:"data"`
	Meta Meta            `json:"meta"`
}

// Meta represents the metadata section of an IonResponse
type Meta struct {
	Version    string    `json:"version"`
	LastUpdate time.Time `json:"last_update"`
	TotalCount int       `json:"total_count"`
	Limit      int       `json:"limit"`
	Offset     int       `json:"offset"`
}

// Pagination represents the necessary elements for a paginated request
type Pagination struct {
	Offset int
	Limit  int
}

// AllItems is a convenience for requesting all items of a given entity
var AllItems = &Pagination{Offset: 0, Limit: -1}

// New takes the credentials required to talk to the API, and the base URL of
// the API and returns a client for talking to the API and an error if any
// issues instantiating the client are encountered
func New(secret string, baseURL string) (*IonClient, error) {
	u, err := url.Parse(baseURL)
	if err != nil {
		return nil, fmt.Errorf("cannot instantiate new ion client: %v", err.Error())
	}

	ic := &IonClient{
		baseURL:     u,
		bearerToken: secret,
		client:      &http.Client{},
	}

	return ic, nil
}

func (ic *IonClient) createURL(endpoint string, params *url.Values, pagination *Pagination) *url.URL {
	u := *ic.baseURL
	u.Path = endpoint

	vals := &url.Values{}
	if params != nil {
		vals = params
	}

	if pagination != nil {
		vals.Set("offset", strconv.Itoa(pagination.Offset))
		vals.Set("limit", strconv.Itoa(pagination.Limit))
	}

	u.RawQuery = vals.Encode()
	return &u
}

func (ic *IonClient) do(method, endpoint string, params *url.Values, payload []byte, headers http.Header, pagination *Pagination) (json.RawMessage, error) {
	if pagination == nil || pagination.Limit > 0 {
		ir, err := ic._do(method, endpoint, params, payload, headers, pagination)
		if err != nil {
			return nil, err
		}

		return ir.Data, nil
	}

	page := &Pagination{Limit: maxPagingLimit, Offset: 0}
	var data json.RawMessage
	data = append(data, []byte("[")...)

	total := 1
	for page.Offset < total {
		ir, err := ic._do(method, endpoint, params, payload, headers, page)
		if err != nil {
			return nil, fmt.Errorf("trouble paging from API: %v", err.Error())
		}
		data = append(data, ir.Data[1:len(ir.Data)-1]...)
		data = append(data, []byte(",")...)
		page.Offset += maxPagingLimit
		total = ir.Meta.TotalCount
	}

	data = append(data[:len(data)-1], []byte("]")...)
	return data, nil
}

func (ic *IonClient) _do(method, endpoint string, params *url.Values, payload []byte, headers http.Header, pagination *Pagination) (*IonResponse, error) {
	u := ic.createURL(endpoint, params, pagination)

	req, err := http.NewRequest(strings.ToUpper(method), u.String(), nil)

	if headers != nil {
		req.Header = headers
	}

	if ic.bearerToken != "" {
		req.Header.Add("Authorization", fmt.Sprintf("Bearer %v", ic.bearerToken))
	}

	resp, err := ic.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed http request: %v", err.Error())
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("bad response from API: %v", resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response body: %v", err.Error())
	}

	var ir IonResponse
	err = json.Unmarshal(body, &ir)
	if err != nil {
		return nil, fmt.Errorf("malformed response: %v", err.Error())
	}

	return &ir, nil
}

func (ic *IonClient) delete(endpoint string, params *url.Values, headers http.Header) (json.RawMessage, error) {
	return ic.do("DELETE", endpoint, params, nil, headers, nil)
}

func (ic *IonClient) get(endpoint string, params *url.Values, headers http.Header, pagination *Pagination) (json.RawMessage, error) {
	return ic.do("GET", endpoint, params, nil, headers, pagination)
}

func (ic *IonClient) post(endpoint string, params *url.Values, payload []byte, headers http.Header) (json.RawMessage, error) {
	return ic.do("POST", endpoint, params, payload, headers, nil)
}

func (ic *IonClient) put(endpoint string, params *url.Values, payload []byte, headers http.Header) (json.RawMessage, error) {
	return ic.do("PUT", endpoint, params, payload, headers, nil)
}
<|endoftext|>


<commit_before>package whois

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"syscall"
	"time"
)

const (
	// DefaultTimeout sets the maximum lifetime of whois requests.
	DefaultTimeout = 30 * time.Second

	// DefaultReadLimit sets the maximum bytes a client will attempt to read from a connection.
	DefaultReadLimit = 1 << 20 // 1 MB
)

// Client represents a whois client. It contains an http.Client, for executing
// some whois Requests.
type Client struct {
	Dial       func(string, string) (net.Conn, error)
	HTTPClient *http.Client
}

// DefaultClient represents a shared whois client with a default timeout, HTTP
// transport, and dialer.
var DefaultClient = NewClient(DefaultTimeout)

// NewClient creates and initializes a new Client with the specified timeout.
func NewClient(timeout time.Duration) *Client {
	dial := func(network, address string) (net.Conn, error) {
		deadline := time.Now().Add(timeout)
		conn, err := net.DialTimeout(network, address, timeout)
		if err != nil {
			return nil, err
		}
		conn.SetDeadline(deadline)
		return conn, nil
	}
	c := &Client{
		Dial:       dial,
		HTTPClient: &http.Client{},
	}
	c.HTTPClient.Transport = &http.Transport{
		Dial:                  c.dial,
		Proxy:                 http.ProxyFromEnvironment,
		TLSHandshakeTimeout:   timeout,
		ResponseHeaderTimeout: timeout,
	}
	return c
}

func (c *Client) dial(network, address string) (net.Conn, error) {
	return c.Dial(network, address)
}

// FetchError reports the underlying error and includes the target host of the fetch operation.
type FetchError struct {
	Err  error
	Host string
}

func (f *FetchError) Error() string {
	return f.Err.Error()
}

// Fetch sends the Request to a whois server.
func (c *Client) Fetch(req *Request) (*Response, error) {
	if req.URL != "" {
		return c.fetchHTTP(req)
	}
	return c.fetchWhois(req)
}

func (c *Client) fetchWhois(req *Request) (*Response, error) {
	conn, err := c.Dial("tcp", req.Host+":43")
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	defer conn.Close()
	if _, err = conn.Write(req.Body); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res := NewResponse(req.Query, req.Host)
	if res.Body, err = ioutil.ReadAll(io.LimitReader(conn, DefaultReadLimit)); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res.DetectContentType("")
	return res, nil
}

func (c *Client) fetchHTTP(req *Request) (*Response, error) {
	hreq, err := httpRequest(req)
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	hres, err := c.HTTPClient.Do(hreq)
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	res := NewResponse(req.Query, req.Host)
	if res.Body, err = ioutil.ReadAll(io.LimitReader(hres.Body, DefaultReadLimit)); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res.DetectContentType(hres.Header.Get("Content-Type"))
	return res, nil
}

func httpRequest(req *Request) (*http.Request, error) {
	var hreq *http.Request
	var err error
	// POST if non-zero Request.Body
	if len(req.Body) > 0 {
		hreq, err = http.NewRequest("POST", req.URL, bytes.NewReader(req.Body))
	} else {
		hreq, err = http.NewRequest("GET", req.URL, nil)
	}
	if err != nil {
		return nil, err
	}
	// Some web whois servers require a Referer header
	hreq.Header.Add("Referer", req.URL)
	return hreq, nil
}

func logError(err error) {
	switch t := err.(type) {
	case syscall.Errno:
		fmt.Fprintf(os.Stderr, "syscall.Errno %d: %s\n", t, err.Error())
	case net.Error:
		fmt.Fprintf(os.Stderr, "net.Error timeout=%t, temp=%t: %s\n", t.Timeout(), t.Temporary(), err.Error())
	default:
		fmt.Fprintf(os.Stderr, "Unknown error %v: %s\n", t, err.Error())
	}
}
<commit_msg>Return an error when missing a whois host.<commit_after>package whois

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"syscall"
	"time"
)

const (
	// DefaultTimeout sets the maximum lifetime of whois requests.
	DefaultTimeout = 30 * time.Second

	// DefaultReadLimit sets the maximum bytes a client will attempt to read from a connection.
	DefaultReadLimit = 1 << 20 // 1 MB
)

// Client represents a whois client. It contains an http.Client, for executing
// some whois Requests.
type Client struct {
	Dial       func(string, string) (net.Conn, error)
	HTTPClient *http.Client
}

// DefaultClient represents a shared whois client with a default timeout, HTTP
// transport, and dialer.
var DefaultClient = NewClient(DefaultTimeout)

// NewClient creates and initializes a new Client with the specified timeout.
func NewClient(timeout time.Duration) *Client {
	dial := func(network, address string) (net.Conn, error) {
		deadline := time.Now().Add(timeout)
		conn, err := net.DialTimeout(network, address, timeout)
		if err != nil {
			return nil, err
		}
		conn.SetDeadline(deadline)
		return conn, nil
	}
	c := &Client{
		Dial:       dial,
		HTTPClient: &http.Client{},
	}
	c.HTTPClient.Transport = &http.Transport{
		Dial:                  c.dial,
		Proxy:                 http.ProxyFromEnvironment,
		TLSHandshakeTimeout:   timeout,
		ResponseHeaderTimeout: timeout,
	}
	return c
}

func (c *Client) dial(network, address string) (net.Conn, error) {
	return c.Dial(network, address)
}

// FetchError reports the underlying error and includes the target host of the fetch operation.
type FetchError struct {
	Err  error
	Host string
}

func (f *FetchError) Error() string {
	return f.Err.Error()
}

// Fetch sends the Request to a whois server.
func (c *Client) Fetch(req *Request) (*Response, error) {
	if req.URL != "" {
		return c.fetchHTTP(req)
	}
	return c.fetchWhois(req)
}

func (c *Client) fetchWhois(req *Request) (*Response, error) {
	if req.Host == "" {
		return nil, &FetchError{fmt.Errorf("no whois host for %s", req.Query), "unknown"}
	}
	conn, err := c.Dial("tcp", req.Host+":43")
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	defer conn.Close()
	if _, err = conn.Write(req.Body); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res := NewResponse(req.Query, req.Host)
	if res.Body, err = ioutil.ReadAll(io.LimitReader(conn, DefaultReadLimit)); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res.DetectContentType("")
	return res, nil
}

func (c *Client) fetchHTTP(req *Request) (*Response, error) {
	hreq, err := httpRequest(req)
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	hres, err := c.HTTPClient.Do(hreq)
	if err != nil {
		return nil, &FetchError{err, req.Host}
	}
	res := NewResponse(req.Query, req.Host)
	if res.Body, err = ioutil.ReadAll(io.LimitReader(hres.Body, DefaultReadLimit)); err != nil {
		logError(err)
		return nil, &FetchError{err, req.Host}
	}
	res.DetectContentType(hres.Header.Get("Content-Type"))
	return res, nil
}

func httpRequest(req *Request) (*http.Request, error) {
	var hreq *http.Request
	var err error
	// POST if non-zero Request.Body
	if len(req.Body) > 0 {
		hreq, err = http.NewRequest("POST", req.URL, bytes.NewReader(req.Body))
	} else {
		hreq, err = http.NewRequest("GET", req.URL, nil)
	}
	if err != nil {
		return nil, err
	}
	// Some web whois servers require a Referer header
	hreq.Header.Add("Referer", req.URL)
	return hreq, nil
}

func logError(err error) {
	switch t := err.(type) {
	case syscall.Errno:
		fmt.Fprintf(os.Stderr, "syscall.Errno %d: %s\n", t, err.Error())
	case net.Error:
		fmt.Fprintf(os.Stderr, "net.Error timeout=%t, temp=%t: %s\n", t.Timeout(), t.Temporary(), err.Error())
	default:
		fmt.Fprintf(os.Stderr, "Unknown error %v: %s\n", t, err.Error())
	}
}
<|endoftext|>