func (tc *TunnelConn) handleIn(fd int) {
targetFd := tc.getPeerFd(fd)
// Block reading while more than 1KB need to be sent
if ts, ok := writeCache[targetFd]; ok && len(ts) > 1024 {
return
}
for num := len(buf); num == len(buf); {
var err error
num, err = tc.readOnce(fd)
if num <= 0 || err != nil {
return
}
toSend := getToSend(targetFd, buf[:num])
sent, left, err := tc.writeOnce(targetFd, toSend)
if err != nil {
return
}
tc.RelTunnel.PassSize += uint64(sent)
if left != nil {
saveLeft(targetFd, left)
continue
}
if _, ok := writeCache[targetFd]; ok {
delete(writeCache, targetFd)
}
if tc.RelTunnel.PassSize > tc.RelTunnel.Limit {
logger.Log(logger.INFO, "Block port [%d]", tc.RelTunnel.EPort)
tc.shutdown()
return
}
}
}
func Run(t *Tunnel) {
var err error
epollFd, err = syscall.EpollCreate(1024)
if err != nil {
logger.Log(logger.ERR, "Create epoll fd error [%s]", err)
os.Exit(-2)
}
for _, step := range initStep {
err = step.Action(t)
if err != nil {
fmt.Fprintf(os.Stderr, step.ErrFmt, err)
os.Exit(-2)
}
}
runTunnel = t
events := make([]syscall.EpollEvent, 10, 10)
for {
en, err := syscall.EpollWait(epollFd, events, 1000)
if err != nil {
logger.Log(logger.ERR, "Wail epoll fd error [%s]", err)
os.Exit(-2)
}
for i := 0; i < en; i++ {
ee := events[i]
if runTunnel.LFd == int(ee.Fd) {
runTunnel.newConn()
continue
}
tc, ok := fdTunnelConn[int(ee.Fd)]
if !ok {
continue
}
if ee.Events&syscall.EPOLLIN != 0 {
tc.handleIn(int(ee.Fd))
}
if ee.Events&syscall.EPOLLOUT != 0 {
tc.handleOut(int(ee.Fd))
}
if ee.Events&syscall.EPOLLHUP != 0 {
tc.shutdown()
}
}
}
}
func Stop() {
logger.Log(logger.INFO, "Stop proxy server")
if runTunnel != nil {
syscall.Close(runTunnel.LFd)
for _, tc := range fdTunnelConn {
syscall.Close(tc.EFd)
syscall.Close(tc.IFd)
}
}
}
func (t *Tunnel) newConn() {
tc := TunnelConn{RelTunnel: t}
for _, step := range pathAddStep {
err := step.Action(&tc)
if err != nil {
logger.Log(logger.ERR, step.ErrFmt, t.EPort, err)
}
}
timePassed := uint(time.Since(t.CheckTime).Seconds())
if timePassed > t.Window {
logger.Log(logger.INFO, "Resume port [%d] Capacity [%d]", t.EPort, t.Limit)
// CheckTime is added on by times of window
t.CheckTime = t.CheckTime.Add(time.Duration(timePassed-timePassed%t.Window) * time.Second)
t.PassSize = 0
} else if t.PassSize > t.Limit {
logger.Log(logger.INFO, "Block port [%d]", t.EPort)
tc.shutdown()
}
fdTunnelConn[tc.IFd], fdTunnelConn[tc.EFd] = &tc, &tc
}
func (tc *TunnelConn) handleOut(fd int) {
out, ok := writeCache[fd]
if !ok {
return
}
num, err := syscall.Write(fd, out)
if err != nil && err != syscall.EAGAIN {
logger.Log(logger.ERR, "Write cache to fd [%d] error [%s]", fd, err)
tc.shutdown()
} else if err == nil && num < len(out) {
writeCache[fd] = out[num:]
} else if err == nil {
delete(writeCache, fd)
}
}
func (tc *TunnelConn) writeOnce(fd int, content []byte) (sent int, left []byte, err error) {
sent, err = syscall.Write(fd, content)
if err != nil && err != syscall.EAGAIN {
logger.Log(logger.ERR, "Write fd [%d] send num [%d] ret [%d] error [%s]", fd, len(content), sent, err)
tc.shutdown()
return
}
if sent > 0 && sent < len(content) {
left = content[sent:]
} else if sent <= 0 && err == syscall.EAGAIN {
sent = 0
left = content
err = nil
}
return
}