This repository has been archived by the owner on Aug 27, 2020. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 31
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
8 changed files
with
410 additions
and
30 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,112 @@ | ||
| package core | ||
|
|
||
| import ( | ||
| "io" | ||
| "math/rand" | ||
| "net" | ||
| "sync" | ||
| "sync/atomic" | ||
| ) | ||
|
|
||
| const ( | ||
| MessageMethodData = iota | ||
| MessageMethodDial | ||
| ) | ||
|
|
||
| type Message struct { | ||
| Method byte | ||
| ConnID uint64 | ||
| MessageID uint64 | ||
| Data []byte | ||
| } | ||
|
|
||
| type MuxConn struct { | ||
| ID uint64 | ||
| muxWS *MuxWebSocket | ||
|
|
||
| mutex sync.Mutex | ||
| buf []byte | ||
| wait chan int | ||
|
|
||
| receiveMessageID uint64 | ||
| sendMessageID *uint64 | ||
| } | ||
|
|
||
| //client use | ||
| func NewMuxConn(muxWS *MuxWebSocket) (conn *MuxConn) { | ||
| return &MuxConn{ | ||
| ID: rand.Uint64(), | ||
| muxWS: muxWS, | ||
| wait: make(chan int), | ||
| sendMessageID: new(uint64), | ||
| } | ||
| } | ||
|
|
||
| func (conn *MuxConn) Write(p []byte) (n int, err error) { | ||
| m := &Message{ | ||
| Method: MessageMethodData, | ||
| ConnID: conn.ID, | ||
| MessageID: conn.SendMessageID(), | ||
| Data: p, | ||
| } | ||
|
|
||
| err = conn.muxWS.SendMessage(m) | ||
| if err != nil { | ||
| return 0, err | ||
| } | ||
| return len(p), nil | ||
| } | ||
|
|
||
| func (conn *MuxConn) Read(p []byte) (n int, err error) { | ||
| if len(conn.buf) == 0 { | ||
| logger.Debugf("%d buf is 0, waiting", conn.ID) | ||
| <-conn.wait | ||
| } | ||
|
|
||
| conn.mutex.Lock() | ||
| logger.Debugf("%d buf: %v", conn.buf) | ||
| n = copy(p, conn.buf) | ||
| conn.buf = conn.buf[n:] | ||
| conn.mutex.Unlock() | ||
| return | ||
| } | ||
|
|
||
| func (conn *MuxConn) HandleMessage(m *Message) (err error) { | ||
| logger.Debugf("handle message %d %d", m.ConnID, m.MessageID) | ||
| for { | ||
| if conn.receiveMessageID == m.MessageID { | ||
| conn.mutex.Lock() | ||
| conn.buf = append(conn.buf, m.Data...) | ||
| conn.receiveMessageID++ | ||
| close(conn.wait) | ||
| conn.wait = make(chan int) | ||
| conn.mutex.Unlock() | ||
| logger.Debugf("handled message %d %d", m.ConnID, m.MessageID) | ||
| return | ||
| } | ||
| <-conn.wait | ||
| } | ||
| return | ||
| } | ||
|
|
||
| func (conn *MuxConn) SendMessageID() (id uint64) { | ||
| id = atomic.LoadUint64(conn.sendMessageID) | ||
| atomic.AddUint64(conn.sendMessageID, 1) | ||
| return | ||
| } | ||
|
|
||
| func (conn *MuxConn) Run(c *net.TCPConn) { | ||
| go func() { | ||
| _, err := io.Copy(c, conn) | ||
| if err != nil { | ||
| logger.Debugf(err.Error()) | ||
| } | ||
| }() | ||
|
|
||
| _, err := io.Copy(conn, c) | ||
| if err != nil { | ||
| logger.Debugf(err.Error()) | ||
| } | ||
|
|
||
| return | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| package core | ||
|
|
||
| import ( | ||
| "net" | ||
| ) | ||
|
|
||
| func (muxWS *MuxWebSocket) ClientListen() { | ||
| for { | ||
| m, err := muxWS.ReceiveMessage() | ||
| if err != nil { | ||
| logger.Debugf(err.Error()) | ||
| return | ||
| } | ||
|
|
||
| //get conn and send message | ||
| conn := muxWS.GetMuxConn(m.ConnID) | ||
| err = conn.HandleMessage(m) | ||
| if err != nil { | ||
| logger.Debugf(err.Error()) | ||
| continue | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (client *Client) OpenMux() (err error) { | ||
| wsConn, _, err := client.Dialer.Dial(client.URL.String(), map[string][]string{ | ||
| "WebSocks-Mux": {"mux"}, | ||
| }) | ||
|
|
||
| if err != nil { | ||
| return | ||
| } | ||
|
|
||
| ws := &WebSocket{ | ||
| conn: wsConn, | ||
| } | ||
|
|
||
| muxWS := NewMuxWebSocket(ws) | ||
| client.MuxWS = muxWS | ||
| return | ||
| } | ||
| func (client *Client) DialMuxConn(host string, conn *net.TCPConn) { | ||
| muxConn := NewMuxConn(client.MuxWS) | ||
|
|
||
| err := muxConn.DialMessage(host) | ||
| if err != nil { | ||
| logger.Errorf(err.Error()) | ||
| err = client.OpenMux() | ||
| if err != nil { | ||
| logger.Errorf(err.Error()) | ||
| } | ||
| return | ||
| } | ||
|
|
||
| muxConn.muxWS.PutMuxConn(muxConn) | ||
|
|
||
| logger.Debugf("dialed mux for %s", host) | ||
|
|
||
| muxConn.Run(conn) | ||
| return | ||
| } | ||
|
|
||
| //client dial remote | ||
| func (conn *MuxConn) DialMessage(host string) (err error) { | ||
| m := &Message{ | ||
| Method: MessageMethodDial, | ||
| MessageID: 18446744073709551615, | ||
| ConnID: conn.ID, | ||
| Data: []byte(host), | ||
| } | ||
|
|
||
| logger.Debugf("dial for %s", host) | ||
|
|
||
| err = conn.muxWS.SendMessage(m) | ||
| if err != nil { | ||
| return | ||
| } | ||
|
|
||
| logger.Debugf("%d %s", conn.ID, host) | ||
| return | ||
| } |
Oops, something went wrong.