1. 增加代码注释

2. 增加通用encoding层

CHANGELOG.md

@@ -199,4 +199,8 @@
 1. 命令行交互切换到cobra库
 2. 新增start、stop、restart、status命令行交互命令
 3. 升级第三方库
-4. 补充代码注释
+4. 补充代码注释
+
+## TODO
+
+1. 补充代码注释

code/client/app/cmd.go

@@ -163,7 +163,7 @@ func (a *App) Status(*cobra.Command, []string) {
 	}
 }
 
-// Vnc handle vnc child process
+// Vnc handle vnc child process handler
 func (a *App) Vnc(*cobra.Command, []string) {
 	defer utils.Recover("vnc.worker")
 

code/client/app/handler.go

@@ -11,6 +11,21 @@ import (
 	"github.com/lwch/natpass/code/network"
 )
 
+/*
+this is the handler function file like http.HandlerFunc.
+
+for linked rule:
+1. get rule from manager
+2. create rule if not exists by its type
+3. create link call NewLink
+4. do the initialize logic for the link
+5. response connect ok message
+6. loop forward
+
+no linked rule:
+TODO
+*/
+
 func (p *program) shellCreate(mgr *rule.Mgr, conn *conn.Conn, msg *network.Msg) {
 	create := msg.GetCreq()
 	tn := mgr.GetLinked(create.GetName(), msg.GetFrom())

code/client/conn/conn.go

@@ -20,14 +20,14 @@ const dropBlockTimeout = 10 * time.Minute
 // Conn connection
 type Conn struct {
 	sync.RWMutex
-	cfg          *global.Configure
-	conn         *network.Conn
+	cfg          *global.Configure            // configure
+	conn         *network.Conn                // connection wrap, read write with timeout
 	read         map[string]chan *network.Msg // link id => channel
 	unknownRead  chan *network.Msg            // read message without link
-	onDisconnect chan string
-	write        chan *network.Msg
-	lockDrop     sync.RWMutex
-	drop         map[string]time.Time
+	onDisconnect chan string                  // on disconnect channel, the value is clientid
+	write        chan *network.Msg            // write queue
+	lockDrop     sync.RWMutex                 // drop mutex
+	drop         map[string]time.Time         // drop cache, drop message when this link is closed
 	// runtime
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -52,11 +52,12 @@ func New(cfg *global.Configure) *Conn {
 	return conn
 }
 
+// connect connect server and write handshake packet
 func (conn *Conn) connect() error {
 	var dial net.Conn
 	var err error
 	if conn.cfg.UseSSL {
-		if conn.cfg.SSLInsecure {
+		if conn.cfg.SSLInsecure { // disable sni
 			rawConn, err := net.Dial("tcp", conn.cfg.Server)
 			if err != nil {
 				logging.Error("raw dial: %v", err)
@@ -96,6 +97,7 @@ func (conn *Conn) close() {
 	}
 }
 
+// writeHandshake send handshake message, default timeout is 5 seconds
 func writeHandshake(conn *network.Conn, cfg *global.Configure) error {
 	var msg network.Msg
 	msg.XType = network.Msg_handshake
@@ -109,6 +111,7 @@ func writeHandshake(conn *network.Conn, cfg *global.Configure) error {
 	return conn.WriteMessage(&msg, 5*time.Second)
 }
 
+// isDrop check the message is dropped by linkid
 func (conn *Conn) isDrop(linkID string) bool {
 	conn.lockDrop.RLock()
 	defer conn.lockDrop.RUnlock()
@@ -116,6 +119,14 @@ func (conn *Conn) isDrop(linkID string) bool {
 	return ok
 }
 
+// addDrop add to drop queue
+func (conn *Conn) addDrop(linkID string) {
+	conn.lockDrop.Lock()
+	defer conn.lockDrop.Unlock()
+	conn.drop[linkID] = time.Now().Add(dropBlockTimeout)
+}
+
+// getChan get read channel by linkid
 func (conn *Conn) getChan(linkID string) chan *network.Msg {
 	conn.RLock()
 	ch := conn.read[linkID]
@@ -126,19 +137,48 @@ func (conn *Conn) getChan(linkID string) chan *network.Msg {
 	return ch
 }
 
-func (conn *Conn) hookDispatch(ch chan *network.Msg, msg *network.Msg) bool {
+// hookDispatch hook message before dispatcher
+func (conn *Conn) hookDispatch(msg *network.Msg) bool {
 	switch msg.GetXType() {
+	// if disconnected add linkid to drop list, and break the handle chain
 	case network.Msg_disconnect:
-		conn.lockDrop.Lock()
-		conn.drop[msg.GetLinkId()] = time.Now().Add(dropBlockTimeout)
-		conn.lockDrop.Unlock()
-		conn.onDisconnect <- msg.GetLinkId()
+		conn.addDrop(msg.GetLinkId())
+		// TODO: no need will block
+		// conn.onDisconnect <- msg.GetLinkId()
 		logging.Info("connection %s disconnected", msg.GetLinkId())
 		return false
 	}
 	return true
 }
 
+// handleLinkedMessage linked message handler, return false means break read loop
+func (conn *Conn) handleLinkedMessage(msg *network.Msg) bool {
+	linkID := msg.GetLinkId()
+	if conn.isDrop(linkID) {
+		return true
+	}
+	if !conn.hookDispatch(msg) {
+		return true
+	}
+	ch := conn.getChan(linkID)
+	select {
+	case ch <- msg:
+	case <-time.After(conn.cfg.WriteTimeout):
+		logging.Error("drop message: %s", msg.GetXType().String())
+		conn.addDrop(linkID)
+	case <-conn.ctx.Done():
+		return false
+	}
+	return true
+}
+
+// handleUnlinkedMessage unlinked message handler, return false means break read loop
+func (conn *Conn) handleUnlinkedMessage(msg *network.Msg) bool {
+	// TODO
+	return true
+}
+
+// loopRead loop read message
 func (conn *Conn) loopRead() {
 	defer utils.Recover("loopRead")
 	defer conn.close()
@@ -146,30 +186,17 @@ func (conn *Conn) loopRead() {
 	var timeout int
 	run := func(msg *network.Msg) bool {
 		timeout = 0
+		// skip keepalive message
 		if msg.GetXType() == network.Msg_keepalive {
 			return true
 		}
 		logging.Debug("read message %s(%s) from %s",
 			msg.GetXType().String(), msg.GetLinkId(), msg.GetFrom())
 		linkID := msg.GetLinkId()
-		if conn.isDrop(linkID) {
-			return true
+		if len(linkID) > 0 {
+			return conn.handleLinkedMessage(msg)
 		}
-		ch := conn.getChan(linkID)
-		if !conn.hookDispatch(ch, msg) {
-			return true
-		}
-		select {
-		case ch <- msg:
-		case <-time.After(conn.cfg.WriteTimeout):
-			logging.Error("drop message: %s", msg.GetXType().String())
-			conn.lockDrop.Lock()
-			conn.drop[msg.GetLinkId()] = time.Now().Add(dropBlockTimeout)
-			conn.lockDrop.Unlock()
-		case <-conn.ctx.Done():
-			return false
-		}
-		return true
+		return conn.handleUnlinkedMessage(msg)
 	}
 	for {
 		msg, _, err := conn.conn.ReadMessage(conn.cfg.ReadTimeout)
@@ -191,6 +218,7 @@ func (conn *Conn) loopRead() {
 	}
 }
 
+// loopWrite loop write message
 func (conn *Conn) loopWrite() {
 	defer utils.Recover("loopWrite")
 	defer conn.close()
@@ -212,6 +240,7 @@ func (conn *Conn) loopWrite() {
 	}
 }
 
+// keepalive loop send keepalive message
 func (conn *Conn) keepalive() {
 	defer utils.Recover("keepalive")
 	defer conn.close()
@@ -237,8 +266,8 @@ func (conn *Conn) AddLink(id string) {
 	conn.Unlock()
 }
 
-// Reset reset message next read
-func (conn *Conn) Reset(id string, msg *network.Msg) {
+// Requeue requeue for next read
+func (conn *Conn) Requeue(id string, msg *network.Msg) {
 	conn.RLock()
 	ch := conn.read[id]
 	conn.RUnlock()
@@ -262,6 +291,7 @@ func (conn *Conn) ChanDisconnect() <-chan string {
 	return conn.onDisconnect
 }
 
+// checkDrop clear timeouted drop queue
 func (conn *Conn) checkDrop() {
 	for {
 		time.Sleep(time.Second)
@@ -298,7 +328,5 @@ func (conn *Conn) ChanClose(id string) {
 	delete(conn.read, id)
 	conn.Unlock()
 
-	conn.lockDrop.Lock()
-	conn.drop[id] = time.Now().Add(dropBlockTimeout)
-	conn.lockDrop.Unlock()
+	conn.addDrop(id)
 }

code/client/rule/code/code.go

@@ -143,7 +143,7 @@ func (code *Code) new(conn *conn.Conn) (string, error) {
 			return "", errWaitingTimeout
 		}
 		if msg.GetXType() != network.Msg_connect_rep {
-			conn.Reset(id, msg)
+			conn.Requeue(id, msg)
 			time.Sleep(code.readTimeout / 10)
 			continue
 		}

code/client/rule/shell/h_new.go

@@ -35,7 +35,7 @@ func (shell *Shell) New(conn *conn.Conn, w http.ResponseWriter, r *http.Request)
 			return
 		}
 		if msg.GetXType() != network.Msg_connect_rep {
-			conn.Reset(id, msg)
+			conn.Requeue(id, msg)
 			time.Sleep(shell.readTimeout / 10)
 			continue
 		}

code/client/rule/vnc/h_new.go

@@ -51,7 +51,7 @@ func (v *VNC) New(conn *conn.Conn, w http.ResponseWriter, r *http.Request) {
 			return
 		}
 		if msg.GetXType() != network.Msg_connect_rep {
-			conn.Reset(id, msg)
+			conn.Requeue(id, msg)
 			time.Sleep(v.readTimeout / 10)
 			continue
 		}

code/network/encoding/encoding.go

@@ -0,0 +1,21 @@
+package encoding
+
+import "io"
+
+// Codec format data to []byte, decode data from []byte
+type Codec interface {
+	// Marshal format data to []byte
+	Marshal(interface{}) ([]byte, error)
+	// Unmarshal decode data from []byte
+	Unmarshal([]byte, interface{}) error
+}
+
+// Compressor compressor interface
+type Compressor interface {
+	// Compress get compress writer
+	Compress(io.Writer) (io.WriteCloser, error)
+	// Decompress get decompress reader
+	Decompress(io.Reader) (io.ReadCloser, error)
+	// SetLevel set compress level
+	SetLevel(int) error
+}

code/network/encoding/gzip/gzip.go

@@ -0,0 +1,88 @@
+package gzip
+
+import (
+	"compress/gzip"
+	"fmt"
+	"io"
+	"sync"
+
+	"github.com/lwch/natpass/code/network/encoding"
+	"github.com/lwch/runtime"
+)
+
+type writer struct {
+	*gzip.Writer
+	pool *sync.Pool
+}
+
+// Close close write and put writer to pool
+func (w *writer) Close() error {
+	w.pool.Put(w)
+	return w.Writer.Close()
+}
+
+type reader struct {
+	*gzip.Reader
+	pool *sync.Pool
+}
+
+// Close close reader and put reader to pool
+func (r *reader) Close() error {
+	r.pool.Put(r)
+	return r.Reader.Close()
+}
+
+type compressor struct {
+	level      int
+	poolWriter [gzip.BestCompression]sync.Pool
+	poolReader sync.Pool
+}
+
+// New create compressor
+func New(level ...int) (encoding.Compressor, error) {
+	if len(level) > 0 {
+		if level[0] < 0 || level[0] > gzip.BestCompression {
+			return nil, fmt.Errorf("invalid gzip compress level: %d", level[0])
+		}
+	} else {
+		level = append(level, 6)
+	}
+	ret := new(compressor)
+	ret.level = level[0]
+	for i := 0; i < gzip.BestCompression; i++ {
+		ret.poolWriter[i].New = func() interface{} {
+			w, err := gzip.NewWriterLevel(io.Discard, i)
+			runtime.Assert(err)
+			return &writer{Writer: w, pool: &ret.poolWriter[i]}
+		}
+	}
+	ret.poolReader.New = func() interface{} {
+		r, err := gzip.NewReader(io.NopCloser(nil))
+		runtime.Assert(err)
+		return &reader{Reader: r, pool: &ret.poolReader}
+	}
+	return ret, nil
+}
+
+// Compress gzip compress
+func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
+	pw := c.poolWriter[c.level].Get().(*writer)
+	pw.Writer.Reset(w)
+	return pw, nil
+}
+
+// Decompress gzip decompress
+func (c *compressor) Decompress(r io.Reader) (io.ReadCloser, error) {
+	pr := c.poolReader.Get().(*reader)
+	pr.Reader.Reset(r)
+	return pr, nil
+}
+
+// SetLevel set compress level
+func (c *compressor) SetLevel(level int) error {
+	if level < 0 || level > gzip.BestCompression {
+		return fmt.Errorf("invalid gzip compress level: %d", level)
+	}
+	c.level = level
+	return nil
+}

code/network/encoding/proto/proto.go

@@ -0,0 +1,33 @@
+package proto
+
+import (
+	"fmt"
+
+	"github.com/lwch/natpass/code/network/encoding"
+	"google.golang.org/protobuf/proto"
+)
+
+type codec struct{}
+
+// New create protobuf codec
+func New() encoding.Codec {
+	return &codec{}
+}
+
+// Marshal protobuf marshal
+func (*codec) Marshal(v interface{}) ([]byte, error) {
+	vv, ok := v.(proto.Message)
+	if !ok {
+		return nil, fmt.Errorf("invalid value type, want proto.Message, got %T", v)
+	}
+	return proto.Marshal(vv)
+}
+
+// Unmarshal protobuf unmarshal
+func (*codec) Unmarshal(data []byte, v interface{}) error {
+	vv, ok := v.(proto.Message)
+	if !ok {
+		return fmt.Errorf("invalid value type, want proto.Message, got %T", v)
+	}
+	return proto.Unmarshal(data, vv)
+}

code/network/network.go

@@ -13,7 +13,8 @@ import (
 	"time"
 
 	"github.com/lwch/logging"
-	"google.golang.org/protobuf/proto"
+	"github.com/lwch/natpass/code/network/encoding"
+	"github.com/lwch/natpass/code/network/encoding/proto"
 )
 
 var errTooLong = errors.New("too long")
@@ -22,12 +23,13 @@ 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
+	c          net.Conn
+	lockRead   sync.Mutex
+	chWrite    chan []byte
+	codec      encoding.Codec
+	compressor encoding.Compressor
+	ctx        context.Context
+	cancel     context.CancelFunc
 }
 
 // NewConn create connection
@@ -36,6 +38,7 @@ func NewConn(c net.Conn) *Conn {
 	conn := &Conn{
 		c:       c,
 		chWrite: make(chan []byte, 1024),
+		codec:   proto.New(),
 		ctx:     ctx,
 		cancel:  cancel,
 	}
@@ -43,66 +46,135 @@ func NewConn(c net.Conn) *Conn {
 	return conn
 }
 
+// SetCompressor set compressor
+func (c *Conn) SetCompressor(cp encoding.Compressor) *Conn {
+	c.compressor = cp
+	return c
+}
+
+// SetCodec set codec
+func (c *Conn) SetCodec(cc encoding.Codec) *Conn {
+	c.codec = cc
+	return c
+}
+
 // Close close connection
 func (c *Conn) Close() {
 	c.c.Close()
 	c.cancel()
 }
 
-func (c *Conn) read(timeout time.Duration) (uint32, uint16, []byte, error) {
+type header struct {
+	Size     uint16
+	Checksum uint32
+}
+
+func (c *Conn) read(timeout time.Duration) ([]byte, error) {
 	c.lockRead.Lock()
 	defer c.lockRead.Unlock()
 	c.c.SetReadDeadline(time.Now().Add(timeout))
-	_, err := io.ReadFull(c.c, c.sizeRead[:])
+	var hdr header
+	err := binary.Read(c.c, binary.BigEndian, &hdr)
 	if err != nil {
-		return 0, 0, nil, err
+		return nil, err
 	}
-	size := binary.BigEndian.Uint16(c.sizeRead[:])
-	enc := binary.BigEndian.Uint32(c.sizeRead[2:])
-	buf := make([]byte, size)
+	buf := make([]byte, hdr.Size)
 	_, err = io.ReadFull(c.c, buf)
 	if err != nil {
-		return 0, 0, nil, err
+		return nil, err
 	}
-	return enc, size, buf, nil
+	if crc32.ChecksumIEEE(buf) != hdr.Checksum {
+		return nil, errChecksum
+	}
+	return buf, nil
+}
+
+func (c *Conn) unserialize(data []byte) (*Msg, error) {
+	if c.compressor != nil {
+		dec, err := c.compressor.Decompress(bytes.NewReader(data))
+		if err != nil {
+			return nil, err
+		}
+		var buffer bytes.Buffer
+		_, err = io.Copy(&buffer, dec)
+		if err != nil {
+			return nil, err
+		}
+		data = buffer.Bytes()
+	}
+	var msg Msg
+	err := c.codec.Unmarshal(data, &msg)
+	if err != nil {
+		return nil, err
+	}
+	return &msg, nil
 }
 
 // ReadMessage read message with timeout
 func (c *Conn) ReadMessage(timeout time.Duration) (*Msg, uint16, error) {
-	enc, size, buf, err := c.read(timeout)
+	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)
+	msg, err := c.unserialize(buf)
 	if err != nil {
 		return nil, 0, err
 	}
-	return &msg, size, nil
+	return msg, uint16(len(buf)), nil
+}
+
+func (c *Conn) serialize(msg *Msg) ([]byte, error) {
+	data, err := c.codec.Marshal(msg)
+	if err != nil {
+		return nil, err
+	}
+	if c.compressor != nil {
+		var buffer bytes.Buffer
+		enc, err := c.compressor.Compress(&buffer)
+		if err != nil {
+			return nil, err
+		}
+		_, err = io.Copy(enc, bytes.NewReader(data))
+		if err != nil {
+			return nil, err
+		}
+		return buffer.Bytes(), nil
+	}
+	return data, nil
+}
+
+func (c *Conn) write(data []byte, timeout time.Duration) error {
+	hdr := header{
+		Size:     uint16(len(data)),
+		Checksum: crc32.ChecksumIEEE(data),
+	}
+	var buffer bytes.Buffer
+	err := binary.Write(&buffer, binary.BigEndian, hdr)
+	if err != nil {
+		return err
+	}
+	_, err = io.Copy(&buffer, bytes.NewReader(data))
+	if err != nil {
+		return err
+	}
+	select {
+	case c.chWrite <- buffer.Bytes():
+		return nil
+	case <-time.After(timeout):
+		return errTimeout
+	}
 }
 
 // WriteMessage write message with timeout
-func (c *Conn) WriteMessage(m *Msg, timeout time.Duration) error {
-	data, err := proto.Marshal(m)
+func (c *Conn) WriteMessage(msg *Msg, timeout time.Duration) error {
+	data, err := c.serialize(msg)
 	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
-	}
+	return c.write(data, timeout)
 }
 
 // RemoteAddr get connection remote address