124 lines
2.4 KiB
Go
124 lines
2.4 KiB
Go
package server
|
|
|
|
import (
|
|
"net/http"
|
|
|
|
"git.pyer.club/kingecg/goemitter"
|
|
"git.pyer.club/kingecg/gologger"
|
|
"git.pyer.club/kingecg/gotunnelserver/util"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type Pipe struct {
|
|
Id string
|
|
Src string
|
|
Dst string
|
|
*goemitter.EventEmitter
|
|
src *websocket.Conn
|
|
dst *websocket.Conn
|
|
stopChan chan int
|
|
logger *gologger.Logger
|
|
}
|
|
|
|
func (p *Pipe) Start() {
|
|
p.stopChan = make(chan int)
|
|
go p.forward(p.src, p.dst)
|
|
go p.forward(p.dst, p.src)
|
|
<-p.stopChan
|
|
p.src.Close()
|
|
p.dst.Close()
|
|
p.Emit("Close")
|
|
}
|
|
func NewPipe(src, dst string) *Pipe {
|
|
return &Pipe{
|
|
Id: util.GenRandomstring(16),
|
|
Src: src,
|
|
Dst: dst,
|
|
EventEmitter: goemitter.NewEmitter(),
|
|
logger: gologger.GetLogger("pipe"),
|
|
}
|
|
|
|
}
|
|
|
|
func (p *Pipe) forward(src, dst *websocket.Conn) {
|
|
for {
|
|
mtype, message, err := src.ReadMessage()
|
|
if err != nil {
|
|
|
|
break
|
|
}
|
|
err = dst.WriteMessage(mtype, message)
|
|
if err != nil {
|
|
|
|
break
|
|
}
|
|
p.logger.Debug("pipe forward:", len(message), "type", mtype)
|
|
}
|
|
p.stopChan <- 1
|
|
}
|
|
|
|
func (s *Server) HandlePipe(conn *websocket.Conn, r *http.Request) {
|
|
session := r.Header.Get("Session")
|
|
if session == "" {
|
|
s.logger.Error("no session header")
|
|
conn.Close()
|
|
return
|
|
}
|
|
cmdSession := s.findSession(session)
|
|
if cmdSession == nil {
|
|
conn.Close()
|
|
s.logger.Error("command session not found")
|
|
return
|
|
}
|
|
lpath := r.URL.Path
|
|
pipeId := lpath[len("/ws/pipe/"):]
|
|
pipe, ok := s.pipes[pipeId]
|
|
|
|
if !ok {
|
|
conn.Close()
|
|
s.logger.Error("pipe not found")
|
|
cmdSession.Send(NewErrorResponse("pipe not found", nil))
|
|
return
|
|
}
|
|
|
|
if pipe.Src == session {
|
|
pipe.src = conn
|
|
} else if pipe.Dst == session {
|
|
pipe.dst = conn
|
|
} else {
|
|
cmdSession.Send(NewErrorResponse("not endpoint of current pipe", nil))
|
|
conn.Close()
|
|
return
|
|
}
|
|
|
|
if pipe.src != nil && pipe.dst != nil {
|
|
go pipe.Start()
|
|
clientCmdSession := s.findSession(pipe.Src)
|
|
clientCmdSession.Send(NcmdConnectionReady(pipe.Id)) // info src endpoint ready and can setup proxy listener
|
|
}
|
|
}
|
|
|
|
func (s *Server) findSession(sidOrName string) *Session {
|
|
clientSession, ok := s.clientSession[sidOrName]
|
|
if ok {
|
|
return clientSession
|
|
}
|
|
agentSession, ok := s.agentSession[sidOrName]
|
|
if ok {
|
|
return agentSession
|
|
}
|
|
|
|
for _, v := range s.clientSession {
|
|
if v.Name == sidOrName {
|
|
return v
|
|
}
|
|
}
|
|
|
|
for _, v := range s.agentSession {
|
|
if v.Name == sidOrName {
|
|
return v
|
|
}
|
|
}
|
|
return nil
|
|
}
|