natpass/code/network/network.go

package network

import (
	"bytes"
	"context"
	"encoding/binary"
	"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")
var errTimeout = errors.New("timeout")

// Conn network connection
type Conn struct {
	c        net.Conn
	lockRead sync.Mutex
	sizeRead [6]byte
	chWrite  chan []byte
	ctx      context.Context
	cancel   context.CancelFunc
}

// NewConn create connection
func NewConn(c net.Conn) *Conn {
	ctx, cancel := context.WithCancel(context.Background())
	conn := &Conn{
		c:       c,
		chWrite: make(chan []byte, 1024),
		ctx:     ctx,
		cancel:  cancel,
	}
	go conn.loopWrite()
	return conn
}

// Close close connection
func (c *Conn) Close() {
	c.c.Close()
	c.cancel()
}

func (c *Conn) read(timeout time.Duration) (uint32, uint16, []byte, 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 0, 0, nil, err
	}
	size := binary.BigEndian.Uint16(c.sizeRead[:])
	enc := binary.BigEndian.Uint32(c.sizeRead[2:])
	buf := make([]byte, size)
	_, err = io.ReadFull(c.c, buf)
	if err != nil {
		return 0, 0, nil, err
	}
	return enc, size, buf, nil
}

// ReadMessage read message with timeout
func (c *Conn) ReadMessage(timeout time.Duration) (*Msg, uint16, error) {
	enc, size, buf, err := c.read(timeout)
	if err != nil {
		return nil, 0, err
	}
	if crc32.ChecksumIEEE(buf) != enc {
		return nil, 0, errChecksum
	}
	var msg Msg
	err = proto.Unmarshal(buf, &msg)
	if err != nil {
		return nil, 0, err
	}
	return &msg, size, nil
}

// WriteMessage write message with timeout
func (c *Conn) WriteMessage(m *Msg, timeout time.Duration) error {
	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))
	copy(buf[len(c.sizeRead):], data)
	select {
	case c.chWrite <- buf:
		return nil
	case <-time.After(timeout):
		return errTimeout
	}
}

// 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()
}

func (c *Conn) loopWrite() {
	defer c.Close()
	for {
		select {
		case <-c.ctx.Done():
			return
		case data := <-c.chWrite:
			_, err := io.Copy(c.c, bytes.NewReader(data))
			if err != nil {
				logging.Error("write data: %v", err)
				return
			}
		}
	}
}