mirror of
https://github.com/lwch/natpass.git
synced 2024-04-21 12:41:54 +00:00
104 lines
2.3 KiB
Go
104 lines
2.3 KiB
Go
package network
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/binary"
|
|
"encoding/hex"
|
|
"errors"
|
|
"hash/crc32"
|
|
"io"
|
|
"math"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/lwch/logging"
|
|
"google.golang.org/protobuf/proto"
|
|
)
|
|
|
|
var errTooLong = errors.New("too long")
|
|
var errChecksum = errors.New("invalid checksum")
|
|
|
|
// Conn network connection
|
|
type Conn struct {
|
|
c net.Conn
|
|
lockRead sync.Mutex
|
|
lockWrite sync.Mutex
|
|
sizeRead [10]byte
|
|
idx uint32
|
|
}
|
|
|
|
// NewConn create connection
|
|
func NewConn(c net.Conn) *Conn {
|
|
return &Conn{c: c}
|
|
}
|
|
|
|
// Close close connection
|
|
func (c *Conn) Close() {
|
|
c.c.Close()
|
|
}
|
|
|
|
// ReadMessage read message with timeout
|
|
func (c *Conn) ReadMessage(timeout time.Duration) (*Msg, error) {
|
|
c.lockRead.Lock()
|
|
defer c.lockRead.Unlock()
|
|
c.c.SetReadDeadline(time.Now().Add(timeout))
|
|
_, err := io.ReadFull(c.c, c.sizeRead[:])
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
size := binary.BigEndian.Uint16(c.sizeRead[:])
|
|
enc := binary.BigEndian.Uint32(c.sizeRead[2:])
|
|
idx := binary.BigEndian.Uint32(c.sizeRead[6:])
|
|
buf := make([]byte, size)
|
|
n, err := io.ReadFull(c.c, buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if crc32.ChecksumIEEE(buf) != enc {
|
|
var msg Msg
|
|
proto.Unmarshal(buf, &msg)
|
|
logging.Info("idx=%d, sum=%d, enc=%d, type=%s, size=%d, buf_size=%d\n%s",
|
|
idx, crc32.ChecksumIEEE(buf), enc, msg.GetXType().String(), size, n, hex.Dump(buf))
|
|
return nil, errChecksum
|
|
}
|
|
var msg Msg
|
|
err = proto.Unmarshal(buf, &msg)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &msg, nil
|
|
}
|
|
|
|
// WriteMessage write message with timeout
|
|
func (c *Conn) WriteMessage(m *Msg, timeout time.Duration) error {
|
|
c.lockWrite.Lock()
|
|
defer c.lockWrite.Unlock()
|
|
data, err := proto.Marshal(m)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(data) > math.MaxUint16 {
|
|
return errTooLong
|
|
}
|
|
buf := make([]byte, len(data)+len(c.sizeRead))
|
|
binary.BigEndian.PutUint16(buf, uint16(len(data)))
|
|
binary.BigEndian.PutUint32(buf[2:], crc32.ChecksumIEEE(data))
|
|
binary.BigEndian.PutUint32(buf[6:], c.idx)
|
|
copy(buf[len(c.sizeRead):], data)
|
|
c.c.SetWriteDeadline(time.Now().Add(timeout))
|
|
_, err = io.Copy(c.c, bytes.NewReader(buf))
|
|
c.idx++
|
|
return err
|
|
}
|
|
|
|
// RemoteAddr get connection remote address
|
|
func (c *Conn) RemoteAddr() net.Addr {
|
|
return c.c.RemoteAddr()
|
|
}
|
|
|
|
// LocalAddr get connection local address
|
|
func (c *Conn) LocalAddr() net.Addr {
|
|
return c.c.LocalAddr()
|
|
}
|