1
0
mirror of https://github.com/schollz/croc.git synced 2024-11-24 16:23:47 +03:00

chore(comm): cleanup comm Write func

+ better handling error.
+ use Send instead of Write since we do not need number of bytes sent.
This commit is contained in:
Khang. Nguyen Khac Nguyen 2020-10-21 15:16:42 +07:00
parent f5624007c8
commit 9f54f2d08f
2 changed files with 12 additions and 12 deletions

View File

@ -77,23 +77,23 @@ func (c *Comm) Close() {
}
}
func (c *Comm) Write(b []byte) (int, error) {
func (c *Comm) Write(b []byte) (n int, err error) {
header := new(bytes.Buffer)
err := binary.Write(header, binary.LittleEndian, uint32(len(b)))
err = binary.Write(header, binary.LittleEndian, uint32(len(b)))
if err != nil {
fmt.Println("binary.Write failed:", err)
}
tmpCopy := append(header.Bytes(), b...)
n, err := c.connection.Write(tmpCopy)
if n != len(tmpCopy) {
if err != nil {
err = fmt.Errorf("wanted to write %d but wrote %d: %w", len(b), n, err)
} else {
err = fmt.Errorf("wanted to write %d but wrote %d", len(b), n)
}
n, err = c.connection.Write(tmpCopy)
if err != nil {
err = fmt.Errorf("connection.Write failed: %w", err)
return
}
// log.Printf("wanted to write %d but wrote %d", n, len(b))
return n, err
if n != len(tmpCopy) {
err = fmt.Errorf("wanted to write %d but wrote %d", len(b), n)
return
}
return
}
func (c *Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {

View File

@ -28,7 +28,7 @@ func Send(c *comm.Comm, key []byte, m Message) (err error) {
if err != nil {
return
}
_, err = c.Write(mSend)
err = c.Send(mSend)
return
}