147 lines
3.3 KiB
Go
147 lines
3.3 KiB
Go
package client
|
|
|
|
import (
|
|
"crypto/tls"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/hashicorp/net-rpc-msgpackrpc"
|
|
"github.com/hashicorp/yamux"
|
|
)
|
|
|
|
const (
|
|
// clientPreamble is the preamble to send before upgrading
|
|
// the connection into a SCADA version 1 connection.
|
|
clientPreamble = "SCADA 1\n"
|
|
|
|
// rpcTimeout is how long of a read deadline we provide
|
|
rpcTimeout = 10 * time.Second
|
|
)
|
|
|
|
// Opts is used to parameterize a Dial
|
|
type Opts struct {
|
|
// Addr is the dial address
|
|
Addr string
|
|
|
|
// TLS controls if TLS is used
|
|
TLS bool
|
|
|
|
// TLSConfig or nil for default
|
|
TLSConfig *tls.Config
|
|
|
|
// Modifies the log output
|
|
LogOutput io.Writer
|
|
}
|
|
|
|
// Client is a SCADA compatible client. This is a bare bones client that
|
|
// only handles the framing and RPC protocol. Higher-level clients should
|
|
// be prefered.
|
|
type Client struct {
|
|
conn net.Conn
|
|
client *yamux.Session
|
|
|
|
closed bool
|
|
closedLock sync.Mutex
|
|
}
|
|
|
|
// Dial is used to establish a new connection over TCP
|
|
func Dial(addr string) (*Client, error) {
|
|
opts := Opts{Addr: addr, TLS: false}
|
|
return DialOpts(&opts)
|
|
}
|
|
|
|
// DialTLS is used to establish a new connection using TLS/TCP
|
|
func DialTLS(addr string, tlsConf *tls.Config) (*Client, error) {
|
|
opts := Opts{Addr: addr, TLS: true, TLSConfig: tlsConf}
|
|
return DialOpts(&opts)
|
|
}
|
|
|
|
// DialOpts is a parameterized Dial
|
|
func DialOpts(opts *Opts) (*Client, error) {
|
|
var conn net.Conn
|
|
var err error
|
|
if opts.TLS {
|
|
conn, err = tls.Dial("tcp", opts.Addr, opts.TLSConfig)
|
|
} else {
|
|
conn, err = net.DialTimeout("tcp", opts.Addr, 10*time.Second)
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return initClient(conn, opts)
|
|
}
|
|
|
|
// initClient does the common initialization
|
|
func initClient(conn net.Conn, opts *Opts) (*Client, error) {
|
|
// Send the preamble
|
|
_, err := conn.Write([]byte(clientPreamble))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("preamble write failed: %v", err)
|
|
}
|
|
|
|
// Wrap the connection in yamux for multiplexing
|
|
ymConf := yamux.DefaultConfig()
|
|
if opts.LogOutput != nil {
|
|
ymConf.LogOutput = opts.LogOutput
|
|
}
|
|
client, _ := yamux.Client(conn, ymConf)
|
|
|
|
// Create the client
|
|
c := &Client{
|
|
conn: conn,
|
|
client: client,
|
|
}
|
|
return c, nil
|
|
}
|
|
|
|
// Close is used to terminate the client connection
|
|
func (c *Client) Close() error {
|
|
c.closedLock.Lock()
|
|
defer c.closedLock.Unlock()
|
|
|
|
if c.closed {
|
|
return nil
|
|
}
|
|
c.closed = true
|
|
c.client.GoAway() // Notify the other side of the close
|
|
return c.client.Close()
|
|
}
|
|
|
|
// RPC is used to perform an RPC
|
|
func (c *Client) RPC(method string, args interface{}, resp interface{}) error {
|
|
// Get a stream
|
|
stream, err := c.Open()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to open stream: %v", err)
|
|
}
|
|
defer stream.Close()
|
|
stream.SetDeadline(time.Now().Add(rpcTimeout))
|
|
|
|
// Create the RPC client
|
|
cc := msgpackrpc.NewCodec(true, true, stream)
|
|
return msgpackrpc.CallWithCodec(cc, method, args, resp)
|
|
}
|
|
|
|
// Accept is used to accept an incoming connection
|
|
func (c *Client) Accept() (net.Conn, error) {
|
|
return c.client.Accept()
|
|
}
|
|
|
|
// Open is used to open an outgoing connection
|
|
func (c *Client) Open() (net.Conn, error) {
|
|
return c.client.Open()
|
|
}
|
|
|
|
// Addr is so that client can act like a net.Listener
|
|
func (c *Client) Addr() net.Addr {
|
|
return c.client.LocalAddr()
|
|
}
|
|
|
|
// NumStreams returns the number of open streams on the client
|
|
func (c *Client) NumStreams() int {
|
|
return c.client.NumStreams()
|
|
}
|