2018-09-23 21:34:29 +02:00
|
|
|
package comm
|
|
|
|
|
|
|
|
import (
|
2018-09-23 21:50:03 +02:00
|
|
|
"bytes"
|
2018-09-23 23:00:31 +02:00
|
|
|
"fmt"
|
2018-09-24 16:51:24 +02:00
|
|
|
"log"
|
2018-09-23 21:34:29 +02:00
|
|
|
"net"
|
2018-09-23 23:15:23 +02:00
|
|
|
"strconv"
|
|
|
|
"strings"
|
2018-09-23 22:32:46 +02:00
|
|
|
"time"
|
2018-09-23 21:34:29 +02:00
|
|
|
)
|
|
|
|
|
|
|
|
// Comm is some basic TCP communication
|
|
|
|
type Comm struct {
|
|
|
|
connection net.Conn
|
|
|
|
}
|
|
|
|
|
|
|
|
// New returns a new comm
|
|
|
|
func New(c net.Conn) Comm {
|
2018-09-23 22:32:46 +02:00
|
|
|
c.SetReadDeadline(time.Now().Add(3 * time.Hour))
|
|
|
|
c.SetDeadline(time.Now().Add(3 * time.Hour))
|
|
|
|
c.SetWriteDeadline(time.Now().Add(3 * time.Hour))
|
2018-09-23 21:34:29 +02:00
|
|
|
return Comm{c}
|
|
|
|
}
|
|
|
|
|
|
|
|
// Connection returns the net.Conn connection
|
|
|
|
func (c Comm) Connection() net.Conn {
|
|
|
|
return c.connection
|
|
|
|
}
|
|
|
|
|
|
|
|
func (c Comm) Write(b []byte) (int, error) {
|
2018-09-23 23:15:23 +02:00
|
|
|
c.connection.Write([]byte(fmt.Sprintf("%0.5d", len(b))))
|
2018-09-23 21:50:03 +02:00
|
|
|
n, err := c.connection.Write(b)
|
2018-09-23 23:00:31 +02:00
|
|
|
if n != len(b) {
|
|
|
|
err = fmt.Errorf("wanted to write %d but wrote %d", n, len(b))
|
|
|
|
}
|
2018-09-23 23:12:45 +02:00
|
|
|
// log.Printf("wanted to write %d but wrote %d", n, len(b))
|
2018-09-23 21:50:03 +02:00
|
|
|
return n, err
|
2018-09-23 21:34:29 +02:00
|
|
|
}
|
|
|
|
|
2018-09-23 23:12:45 +02:00
|
|
|
func (c Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
|
2018-09-23 23:15:23 +02:00
|
|
|
// read until we get 5 bytes
|
2018-09-24 16:51:24 +02:00
|
|
|
tmp := make([]byte, 5)
|
|
|
|
n, err := c.connection.Read(tmp)
|
2018-09-23 21:50:03 +02:00
|
|
|
if err != nil {
|
|
|
|
return
|
|
|
|
}
|
2018-09-24 16:51:24 +02:00
|
|
|
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)
|
2018-09-23 23:12:45 +02:00
|
|
|
for {
|
2018-09-24 16:51:24 +02:00
|
|
|
// see if we have enough bytes
|
2018-09-24 04:24:57 +02:00
|
|
|
bs = bytes.Trim(bs, "\x00")
|
2018-09-23 23:15:23 +02:00
|
|
|
if len(bs) == 5 {
|
2018-09-23 23:12:45 +02:00
|
|
|
break
|
|
|
|
}
|
2018-09-24 16:51:24 +02:00
|
|
|
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...)
|
2018-09-23 23:15:23 +02:00
|
|
|
}
|
2018-09-24 16:51:24 +02:00
|
|
|
|
2018-09-23 23:15:23 +02:00
|
|
|
numBytes, err = strconv.Atoi(strings.TrimLeft(string(bs), "0"))
|
|
|
|
if err != nil {
|
|
|
|
return nil, 0, nil, err
|
2018-09-23 23:12:45 +02:00
|
|
|
}
|
2018-09-24 04:24:57 +02:00
|
|
|
buf = make([]byte, numBytes)
|
|
|
|
tmp = make([]byte, numBytes)
|
|
|
|
bufStart := 0
|
2018-09-23 21:50:03 +02:00
|
|
|
for {
|
2018-09-24 16:51:24 +02:00
|
|
|
n, err := c.connection.Read(tmp)
|
2018-09-23 21:50:03 +02:00
|
|
|
if err != nil {
|
2018-09-24 16:51:24 +02:00
|
|
|
return nil, 0, nil, err
|
2018-09-23 21:50:03 +02:00
|
|
|
}
|
2018-09-24 16:51:24 +02:00
|
|
|
tmpCopy := make([]byte, n)
|
|
|
|
// Copy the buffer so it doesn't get changed while read by the recipient.
|
|
|
|
copy(tmpCopy, tmp[:n])
|
|
|
|
|
|
|
|
tmpCopy = bytes.TrimSpace(tmpCopy)
|
|
|
|
tmpCopy = bytes.Replace(tmpCopy, []byte(" "), []byte{}, -1)
|
|
|
|
tmpCopy = bytes.Trim(tmpCopy, "\x00")
|
|
|
|
copy(buf[bufStart:bufStart+len(tmpCopy)], tmpCopy[:])
|
2018-09-24 15:10:31 +02:00
|
|
|
bufStart = len(buf)
|
2018-09-24 04:24:57 +02:00
|
|
|
if bufStart < numBytes {
|
2018-09-24 16:51:24 +02:00
|
|
|
// shrink the amount we need to read
|
2018-09-24 04:24:57 +02:00
|
|
|
tmp = tmp[:numBytes-bufStart]
|
2018-09-23 21:50:03 +02:00
|
|
|
} else {
|
|
|
|
break
|
|
|
|
}
|
|
|
|
}
|
2018-09-24 16:51:24 +02:00
|
|
|
log.Printf("wanted %d and got %d", numBytes, len(buf))
|
2018-09-23 21:34:29 +02:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
// Send a message
|
|
|
|
func (c Comm) Send(message string) (err error) {
|
2018-09-23 21:50:03 +02:00
|
|
|
_, err = c.Write([]byte(message))
|
2018-09-23 21:34:29 +02:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
// Receive a message
|
|
|
|
func (c Comm) Receive() (s string, err error) {
|
2018-09-23 23:12:45 +02:00
|
|
|
b, _, _, err := c.Read()
|
2018-09-23 21:50:03 +02:00
|
|
|
s = string(b)
|
2018-09-23 21:34:29 +02:00
|
|
|
return
|
|
|
|
}
|