127 lines
2.7 KiB
Go
127 lines
2.7 KiB
Go
package comm
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"fmt"
|
|
"net"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// Comm is some basic TCP communication
|
|
type Comm struct {
|
|
connection net.Conn
|
|
writer *bufio.Writer
|
|
reader *bufio.Reader
|
|
}
|
|
|
|
// New returns a new comm
|
|
func New(n net.Conn) *Comm {
|
|
c := new(Comm)
|
|
c.connection = n
|
|
c.connection.SetReadDeadline(time.Now().Add(3 * time.Hour))
|
|
c.connection.SetDeadline(time.Now().Add(3 * time.Hour))
|
|
c.connection.SetWriteDeadline(time.Now().Add(3 * time.Hour))
|
|
c.writer = bufio.NewWriter(n)
|
|
// c.connection = bufio.NewReader(n)
|
|
return c
|
|
}
|
|
|
|
// Connection returns the net.TCPConn connection
|
|
func (c *Comm) Connection() net.Conn {
|
|
return c.connection
|
|
}
|
|
|
|
// Close closes the connection
|
|
func (c *Comm) Close() {
|
|
c.connection.Close()
|
|
}
|
|
|
|
func (c *Comm) Write(b []byte) (int, error) {
|
|
c.writer.Write([]byte(fmt.Sprintf("%0.6d", len(b))))
|
|
c.writer.Flush()
|
|
n, err := c.writer.Write(b)
|
|
if n != len(b) {
|
|
err = fmt.Errorf("wanted to write %d but wrote %d", n, len(b))
|
|
}
|
|
if err == nil {
|
|
c.writer.Flush()
|
|
}
|
|
// log.Printf("wanted to write %d but wrote %d", n, len(b))
|
|
return n, err
|
|
}
|
|
|
|
// func (c *Comm) Flush() {
|
|
// c.connection.Flush()
|
|
// }
|
|
|
|
func (c *Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
|
|
// read until we get 6 bytes
|
|
tmp := make([]byte, 6)
|
|
n, err := c.connection.Read(tmp)
|
|
if err != nil {
|
|
return
|
|
}
|
|
tmpCopy := make([]byte, n)
|
|
// Copy the buffer so it doesn't get changed while read by the recipient.
|
|
copy(tmpCopy, tmp[:n])
|
|
bs = tmpCopy
|
|
|
|
tmp = make([]byte, 1)
|
|
for {
|
|
// see if we have enough bytes
|
|
bs = bytes.Trim(bs, "\x00")
|
|
if len(bs) == 6 {
|
|
break
|
|
}
|
|
n, err := c.connection.Read(tmp)
|
|
if err != nil {
|
|
return nil, 0, nil, err
|
|
}
|
|
tmpCopy = make([]byte, n)
|
|
// Copy the buffer so it doesn't get changed while read by the recipient.
|
|
copy(tmpCopy, tmp[:n])
|
|
bs = append(bs, tmpCopy...)
|
|
}
|
|
|
|
numBytes, err = strconv.Atoi(strings.TrimLeft(string(bs), "0"))
|
|
if err != nil {
|
|
return nil, 0, nil, err
|
|
}
|
|
buf = []byte{}
|
|
tmp = make([]byte, numBytes)
|
|
for {
|
|
n, err := c.connection.Read(tmp)
|
|
if err != nil {
|
|
return nil, 0, nil, err
|
|
}
|
|
tmpCopy = make([]byte, n)
|
|
// Copy the buffer so it doesn't get changed while read by the recipient.
|
|
copy(tmpCopy, tmp[:n])
|
|
buf = append(buf, bytes.TrimRight(tmpCopy, "\x00")...)
|
|
if len(buf) < numBytes {
|
|
// shrink the amount we need to read
|
|
tmp = tmp[:numBytes-len(buf)]
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
// log.Printf("wanted %d and got %d", numBytes, len(buf))
|
|
return
|
|
}
|
|
|
|
// Send a message
|
|
func (c *Comm) Send(message string) (err error) {
|
|
_, err = c.Write([]byte(message))
|
|
return
|
|
}
|
|
|
|
// Receive a message
|
|
func (c *Comm) Receive() (s string, err error) {
|
|
b, _, _, err := c.Read()
|
|
s = string(b)
|
|
return
|
|
}
|