diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e08aec..83e49d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -199,4 +199,8 @@ 1. 命令行交互切换到cobra库 2. 新增start、stop、restart、status命令行交互命令 3. 升级第三方库 -4. 补充代码注释 \ No newline at end of file +4. 补充代码注释 + +## TODO + +1. 补充代码注释 \ No newline at end of file diff --git a/code/client/app/cmd.go b/code/client/app/cmd.go index 031e8e8..566f406 100644 --- a/code/client/app/cmd.go +++ b/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") diff --git a/code/client/app/callback.go b/code/client/app/handler.go similarity index 92% rename from code/client/app/callback.go rename to code/client/app/handler.go index 2de8177..31804f1 100644 --- a/code/client/app/callback.go +++ b/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()) diff --git a/code/client/conn/conn.go b/code/client/conn/conn.go index a1c167c..b6b3bdb 100644 --- a/code/client/conn/conn.go +++ b/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) } diff --git a/code/client/rule/code/code.go b/code/client/rule/code/code.go index 5016d04..fc62eab 100644 --- a/code/client/rule/code/code.go +++ b/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 } diff --git a/code/client/rule/shell/h_new.go b/code/client/rule/shell/h_new.go index b2f8495..1f8394a 100644 --- a/code/client/rule/shell/h_new.go +++ b/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 } diff --git a/code/client/rule/vnc/h_new.go b/code/client/rule/vnc/h_new.go index 6d30a41..959833a 100644 --- a/code/client/rule/vnc/h_new.go +++ b/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 } diff --git a/code/network/encoding/encoding.go b/code/network/encoding/encoding.go new file mode 100644 index 0000000..20abca4 --- /dev/null +++ b/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 +} diff --git a/code/network/encoding/gzip/gzip.go b/code/network/encoding/gzip/gzip.go new file mode 100644 index 0000000..4501319 --- /dev/null +++ b/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 +} diff --git a/code/network/encoding/proto/proto.go b/code/network/encoding/proto/proto.go new file mode 100644 index 0000000..0e836f1 --- /dev/null +++ b/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) +} diff --git a/code/network/network.go b/code/network/network.go index 9130ecd..4627107 100644 --- a/code/network/network.go +++ b/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