Allow queries over arbitrary net.Conn

This introduces NewLivestatusFromDialer that takes a custom dial
function.

livestatus.go

@@ -1,10 +1,13 @@
 // Package livestatus provides a binding to MK Livestatus sockets.
 package livestatus
 
+import "net"
+
 // Livestatus is a binding instance.
 type Livestatus struct {
 	network string
 	address string
+	dialer  func() (net.Conn, error)
 }
 
 // Query creates a new query instance on a spacific table.
@@ -19,3 +22,11 @@ func NewLivestatus(network, address string) *Livestatus {
 		address: address,
 	}
 }
+
+// NewLivestatusWithDialer creates a new binding that uses the net.Conn returned
+// by the provided dialer function.
+func NewLivestatusWithDialer(dialer func() (net.Conn, error)) *Livestatus {
+	return &Livestatus{
+		dialer: dialer,
+	}
+}

query.go

@@ -82,14 +82,20 @@ func (q *Query) Exec() (*Response, error) {
 	for {
 		data = make([]byte, 1024)
 
-		_, err = conn.Read(data)
+		n, err := conn.Read(data)
 		if err == io.EOF {
 			break
 		} else if err != nil {
 			return nil, err
 		}
 
 		buf.Write(bytes.TrimRight(data, "\x00"))
+
+		// New line signals the end of content. This check helps
+		// if the connection is not forcibly closed
+		if data[n-1] == byte('\n') {
+			break
+		}
 	}
 
 	if buf.Len() == 0 {
@@ -122,7 +128,11 @@ func (q *Query) buildCmd() string {
 }
 
 func (q *Query) dial() (net.Conn, error) {
-	return net.Dial(q.ls.network, q.ls.address)
+	if q.ls.dialer != nil {
+		return q.ls.dialer()
+	} else {
+		return net.Dial(q.ls.network, q.ls.address)
+	}
 }
 
 func (q *Query) parse(data []byte) ([]Record, error) {