Merge branch 'transfer' into 'master'

Transfer

See merge request !13

cli.go

@@ -46,7 +46,7 @@ func init() {
 
 func initial() {
 	if global.Config.PprofPort > 0 {
-		safe.SafeGo(func() {
+		safe.Go(func() {
 			addr := fmt.Sprintf("127.0.0.1:%d", global.Config.PprofPort)
 			log.Run().Infof("enable pprof: %s", addr)
 			log.Run().Error(http.ListenAndServe(addr, nil))

client/client.go

@@ -17,8 +17,7 @@ type Client struct {
 	localAddr     net.Addr
 	natTunnel     getty.Client
 	session       getty.Session
-	connStore    sync.Map
-	sessionStore sync.Map
+	transferStore sync.Map
 	dnsCache *freecache.Cache
 }
 
@@ -39,8 +38,7 @@ func NewCli(cfg *model.Configure, host string, port int) *Client {
 		host:          host,
 		port:          port,
 		localAddr:     localAddr,
-		connStore:    sync.Map{},
-		sessionStore: sync.Map{},
+		transferStore: sync.Map{},
 		dnsCache:      freecache.NewCache(1024),
 	}
 
@@ -50,4 +48,3 @@ func NewCli(cfg *model.Configure, host string, port int) *Client {
 func (client *Client) StartUp() {
 	client.connect()
 }
-

client/cluster.go

@@ -23,7 +23,7 @@ type ClusterClient struct {
 
 func (c *ClusterClient) Start() {
 	c.check()
-	safe.SafeGo(func() {
+	safe.Go(func() {
 		var timer = time.NewTimer(5 * time.Minute)
 		for {
 			c.connectNatServers()
@@ -32,7 +32,7 @@ func (c *ClusterClient) Start() {
 		}
 	})
 	if global.Config.Redial.Valid() {
-		safe.SafeGo(func() {
+		safe.Go(func() {
 			// 加上随机 防止vps在同时间重启
 			duration := c.randomDuration()
 			log.Run().Infof("Redial interval %+v", duration)
@@ -142,7 +142,7 @@ func (c *ClusterClient) check() {
 		url = global.Config.NetCheckUrl
 	}
 
-	safe.SafeGo(func() {
+	safe.Go(func() {
 		var timer = time.NewTimer(interval)
 		for {
 			timer.Reset(interval)

client/getty.go

@@ -75,7 +75,7 @@ func (client *Client) Redial(tag string) {
 	if !client.config.Redial.Valid() {
 		return
 	}
-	log.Run().Infof("[Redial %s] Send offline message", tag)
+	log.Run().Infof("[Redial %s] seed offline message", tag)
 	if _, _, err := client.session.WritePkg(OfflinePacket, 0); err != nil {
 		log.Run().Errorf("[Redial %s] write offline to server error %s", tag, err.Error())
 	}

client/transfer.go

+package client
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"sync"
+	"time"
+	"virjar.com/majora-go/common"
+	"virjar.com/majora-go/log"
+	"virjar.com/majora-go/protocol"
+	"virjar.com/majora-go/safe"
+	"virjar.com/majora-go/trace"
+)
+
+type Transfer struct {
+	serialNumber             int64
+	client                   *Client
+	upstreamConn             *net.TCPConn
+	recorder                 trace.Recorder
+	transferChan             chan *protocol.MajoraPacket
+	transferToUpstreamFunc   func(t *Transfer, p *protocol.MajoraPacket)
+	transferToDownstreamFunc func(t *Transfer, data []byte, err error)
+	once                     sync.Once
+	cancel                   chan struct{}
+}
+
+func NewTransfer(serialNumber int64, client *Client, conn *net.TCPConn, recorder trace.Recorder) *Transfer {
+	return &Transfer{
+		serialNumber: serialNumber,
+		client:       client,
+		upstreamConn: conn,
+		recorder:     recorder,
+		transferChan: make(chan *protocol.MajoraPacket, 10),
+		once:         sync.Once{},
+		cancel:       make(chan struct{}, 0),
+	}
+}
+
+func (t *Transfer) SetTransferToUpstreamFunc(f func(t *Transfer, p *protocol.MajoraPacket)) {
+	t.transferToUpstreamFunc = f
+}
+
+func (t *Transfer) SetTransferToDownstreamFunc(f func(t *Transfer, data []byte, err error)) {
+	t.transferToDownstreamFunc = f
+}
+
+func (t *Transfer) TransferToUpstream(p *protocol.MajoraPacket) {
+	t.transferChan <- p
+}
+
+func (t *Transfer) Start() {
+	if t.transferToUpstreamFunc == nil {
+		panic(errors.New("transferToUpstreamFunc is nil"))
+	}
+
+	if t.transferToDownstreamFunc == nil {
+		panic(errors.New("transferToDownstreamFunc is nil"))
+	}
+
+	safe.Go(func() {
+		for {
+			select {
+			case p := <-t.transferChan:
+				t.transferToUpstreamFunc(t, p)
+			case <-t.cancel:
+				return
+			}
+		}
+	})
+
+	safe.Go(func() {
+		traceRecorder := t.recorder
+		traceRecorder.RecordEvent(trace.UpStreamEvent, fmt.Sprintf("Ready read from upstream (sn:%d)", t.serialNumber))
+		log.Run().Debugf("[handleUpStream] %d-> handleUpStream start...", t.serialNumber)
+		for {
+			buf := make([]byte, common.BufSize)
+			cnt, err := t.upstreamConn.Read(buf)
+			t.transferToDownstreamFunc(t, buf[0:cnt], err)
+			if err != nil {
+				break
+			}
+		}
+	})
+}
+
+func (t *Transfer) Close() {
+	t.once.Do(func() {
+		readDeadLine := time.Now().Add(3 * time.Millisecond)
+		t.recorder.RecordEvent(trace.DisconnectEvent, fmt.Sprintf("Set upstream read deadline:%s (sn:%d)",
+			readDeadLine.Format("2006-01-02 15:04:05.000000"), t.serialNumber))
+		err := t.upstreamConn.SetReadDeadline(readDeadLine)
+		if err != nil {
+			t.recorder.RecordErrorEvent(trace.DisconnectEvent,
+				fmt.Sprintf("Set upstream read deadline failed (sn:%d)", t.serialNumber), err)
+			_ = t.upstreamConn.Close()
+		}
+		close(t.cancel)
+	})
+}

conf/majora-dev.yaml

-tunnel_addr: majora-vps-zj.virjar.com:5879
+tunnel_addr: 127.0.0.1:5879
 dns_server: 114.114.114.114:53
-daemon: true
+#daemon: true
 log_level: debug
 log_path: ./majora-log/
 reconn_intervalz: 5s
@@ -8,11 +8,11 @@ net_check_interval: 5s
 dns_cache_duration: 10m
 net_check_url: https://www.baidu.com[extra]
 
-redial:
-  command: /bin/bash
-  exec_path: /root/ppp_redial.sh
-  redial_duration: 5m
-  wait_time: 10s
+#redial:
+#  command: /bin/bash
+#  exec_path: /root/ppp_redial.sh
+#  redial_duration: 5m
+#  wait_time: 10s
 
 extra:
   account: superman

safe/safe.go

@@ -2,7 +2,7 @@ package safe
 
 import "virjar.com/majora-go/log"
 
-func SafeGo(f func()) {
+func Go(f func()) {
 	go func() {
 		defer func() {
 			if err := recover(); err != nil {

trace/trace.go

@@ -22,7 +22,7 @@ var (
 )
 
 func init() {
-	safe.SafeGo(func() {
+	safe.Go(func() {
 		for {
 			e := <-sessionEventChan
 			if e.Err != nil {
@@ -133,15 +133,9 @@ func acquireRecorder(sessionId string, host, user string, enable bool) Recorder
 
 }
 
-type Session struct {
-	Recorder Recorder
-}
-
-func NewSession(sessionId string, host string, user string, enable bool) *Session {
+func NewRecorder(sessionId string, host string, user string, enable bool) Recorder {
 	if len(sessionId) == 0 {
 		sessionId = sessionIdNop
 	}
-	return &Session{
-		Recorder: acquireRecorder(sessionId, host, user, enable),
-	}
+	return acquireRecorder(sessionId, host, user, enable)
 }